WIP: AI Agent v0

app/.env.default

@@ -7,3 +7,4 @@ VLLM_LLAMA31_ENV_FILE=""
 JWT_SECRET=test-secret-456
 DJANGO_SECRET_KEY=django-insecure-default
 HF_TOKEN=<your-hf-token> # Get this from Hugging Face
+TAVILY_API_KEY=<your-tavily-api-key>

app/api/agent_control/Dockerfile

@@ -0,0 +1,46 @@
+# install ubutntu base image 
+FROM ubuntu:20.04
+ENV TZ=America/Los_Angeles
+ARG DEBIAN_FRONTEND=noninteractive
+
+# Update the package repository and install some default tools
+RUN apt-get update && apt-get install -y \
+    vim \
+    nano \
+    software-properties-common  \ 
+    git \
+    htop \
+    screen \
+    tmux \
+    unzip \
+    zip \
+    curl \
+    wget 
+
+# add deadsnakes for newer python versions
+RUN add-apt-repository ppa:deadsnakes/ppa -y && apt-get update
+# Install the specific version of Python 
+RUN apt-get install -y python3.11 python3.11-venv python3.11-dev
+
+# Set Python3.11 as the default Python3
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1
+
+# Ensure pip is installed and upgrade it
+RUN python3 -m ensurepip --upgrade && \
+    python3 -m pip install --upgrade pip setuptools wheel
+
+# Verify the Python version
+RUN python3 --version
+
+COPY requirements_agent_env.txt . 
+# install python dependencies 
+RUN pip install --no-cache-dir -r requirements_agent_env.txt
+
+# Set the working directory
+WORKDIR /app
+
+# Copy files (optional)
+COPY . /app
+
+# Command to run when the container starts
+CMD ["/bin/bash"]

app/api/agent_control/README.md

@@ -0,0 +1,9 @@
+```bash
+python -m venv agent_env
+source agent_env/bin/activate
+pip install -r requirements_agent_env.txt
+```
+
+```
+python agent.py
+```

app/api/agent_control/agent.py

@@ -0,0 +1,58 @@
+from custom_llm import CustomLLM
+from utils import poll_requests, setup_executer
+from code_tool import CodeInterpreterFunctionTool
+from langchain.memory import ConversationBufferMemory
+from langchain_community.tools.tavily_search import TavilySearchResults
+import asyncio
+import os 
+import jwt
+import json
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from fastapi.responses import StreamingResponse
+
+
+
+app = FastAPI()
+json_payload = json.loads('{"team_id": "tenstorrent", "token_id":"debug-test"}')
+jwt_secret = os.getenv("JWT_SECRET")
+encoded_jwt = jwt.encode(json_payload, jwt_secret, algorithm="HS256")
+
+class RequestPayload(BaseModel):
+    message: str
+    thread_id: str
+
+
+llm_container_name = os.getenv("LLM_CONTAINER_NAME")
+llm = CustomLLM(server_url=f"http://{llm_container_name}:7000/v1/chat/completions", encoded_jwt=encoded_jwt, 
+                streaming=True)
+memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+os.environ["TAVILY_API_KEY"] = os.getenv("TAVILY_API_KEY")
+
+search = TavilySearchResults(
+    max_results=2,
+    include_answer=True,
+    include_raw_content=True)
+
+# TODO: enable code agent
+# os.environ["E2B_API_KEY"] = os.getenv("E2B_API_KEY")
+# code_interpreter = CodeInterpreterFunctionTool()
+# code_interpreter_tool = code_interpreter.to_langchain_tool()
+tools = [search]
+agent_executer = setup_executer(llm, memory, tools)
+config = {"configurable": {"thread_id": "abc-123"}}
+# asyncio.run(poll_requests(agent_executer, config, tools, memory)) # TODO: enable to run without server 
+
+@app.post("/poll_requests")
+async def handle_requests(payload: RequestPayload):
+    config = {"configurable": {"thread_id": payload.thread_id}}
+    try:
+        # use await to prevent handle_requests from blocking, allow other tasks to execute
+        return StreamingResponse(poll_requests(agent_executer, config, tools, memory, payload.message), media_type="text/plain")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# health check
+@app.get("/")
+def read_root():
+    return {"message": "Server is running"}

app/api/agent_control/code_tool.py

@@ -0,0 +1,106 @@
+from pydantic import BaseModel, Field
+from typing import List, Sequence, Tuple, Any
+from langchain_core.messages import BaseMessage
+from langchain.agents.output_parsers.tools import ToolAgentAction
+from e2b_code_interpreter import Sandbox
+from langchain_core.tools import Tool
+from langchain_core.messages import ToolMessage
+
+
+import os
+import json 
+
+def format_to_tool_messages(
+    intermediate_steps: Sequence[Tuple[ToolAgentAction, dict]],
+) -> List[BaseMessage]:
+    messages = []
+    for agent_action, observation in intermediate_steps:
+        if agent_action.tool == CodeInterpreterFunctionTool.tool_name:
+            new_messages = CodeInterpreterFunctionTool.format_to_tool_message(
+                agent_action,
+                observation,
+            )
+            messages.extend([new for new in new_messages if new not in messages])
+        else:
+            # Handle other tools
+            print("Not handling tool: ", agent_action.tool)
+
+    return messages
+
+class LangchainCodeInterpreterToolInput(BaseModel):
+    code: str = Field(description="Python code to execute.")
+
+
+class CodeInterpreterFunctionTool:
+    """
+    This class calls arbitrary code against a Python Jupyter notebook.
+    It requires an E2B_API_KEY to create a sandbox.
+    """
+
+    tool_name: str = "code_interpreter"
+
+    def __init__(self):
+        # Instantiate the E2B sandbox - this is a long lived object
+        # that's pinging E2B cloud to keep the sandbox alive.
+        if "E2B_API_KEY" not in os.environ:
+            raise Exception(
+                "Code Interpreter tool called while E2B_API_KEY environment variable is not set. Please get your E2B api key here https://e2b.dev/docs and set the E2B_API_KEY environment variable."
+            )
+        self.code_interpreter = Sandbox(timeout=1800)
+
+    def call(self, parameters: dict, **kwargs: Any):
+        code = parameters.get("code", "")
+        if code.startswith("```"):
+            code = code[3:]
+        if code.endswith("[DONE]"):
+            # TODO: check if this needs to be parsed 
+            pass 
+        if code.endswith("```"):
+            code = code[:-3]
+        elif code.endswith("```\n"):
+            code = code[:-4]
+        print(f"***Code Interpreting...\n{code}\n====")
+        execution = self.code_interpreter.run_code(code)
+        return {
+            "results": execution.results,
+            "stdout": execution.logs.stdout,
+            "stderr": execution.logs.stderr,
+            "error": execution.error,
+        }
+
+    def close(self):
+        self.code_interpreter.kill()
+
+    # langchain does not return a dict as a parameter, only a code string
+    def langchain_call(self, code: str):
+        return self.call({"code": code})
+
+    def to_langchain_tool(self) -> Tool:
+        tool = Tool(
+            name=self.tool_name,
+            description="Execute python code in a Jupyter notebook cell and returns any rich data (eg charts), stdout, stderr, and error.",
+            func=self.langchain_call,
+        )
+        tool.args_schema = LangchainCodeInterpreterToolInput
+        return tool
+
+    @staticmethod
+    def format_to_tool_message(
+        agent_action: ToolAgentAction,
+        observation: dict,
+    ) -> List[BaseMessage]:
+        """
+        Format the output of the CodeInterpreter tool to be returned as a ToolMessage.
+        """
+        new_messages = list(agent_action.message_log)
+
+        # TODO: Add info about the results for the LLM
+        content = json.dumps(
+            {k: v for k, v in observation.items() if k not in ("results")}, indent=2
+        )
+        print(observation, agent_action, content)
+        new_messages.append(
+            ToolMessage(content=content, tool_call_id=agent_action.tool_call_id)
+        )
+
+        return new_messages