Implement OracleBot class and update agent configuration

Dockerfile

@@ -0,0 +1,16 @@
+FROM python:3.13-slim-bookworm AS base
+FROM base AS builder
+COPY --from=ghcr.io/astral-sh/uv:0.7.13 /uv /bin/uv
+ENV UV_COMPILE_BYTECODE=1 UV_LINK_MODE=copy
+WORKDIR /app
+COPY uv.lock pyproject.toml /app/
+RUN --mount=type=cache,target=/root/.cache/uv \
+  uv sync --frozen --no-install-project --no-dev
+COPY . /app
+RUN --mount=type=cache,target=/root/.cache/uv \
+  uv sync --frozen --no-dev
+FROM base
+COPY --from=builder /app /app
+ENV PATH="/app/.venv/bin:$PATH"
+WORKDIR /app
+CMD ["python", "/app/main.py"]

agent/__init__.py

@@ -1,11 +1,11 @@
-from .agent import get_agent
+from .agent import OracleBot
 from .config import create_agent_config
 from .prompts import (
     DEFAULT_PROMPT
 )
 
 __all__ = [
-    "get_agent", 
+    "OracleBot",
     "create_agent_config", 
     "DEFAULT_PROMPT"
 ]

agent/agent.py

@@ -3,46 +3,64 @@ from typing_extensions import TypedDict
 from langgraph.graph import StateGraph, START, END
 from langgraph.prebuilt import tools_condition
 from agent.nodes import call_model, tool_node
-from agent.state import State
+from langgraph.graph import MessagesState
 from langchain_core.messages import AIMessage, HumanMessage
+from langgraph.checkpoint.memory import InMemorySaver
+from agent.config import create_agent_config
 
+class OracleBot:
+    def __init__(self):
+        print("Initializing OracleBot")
+        self.name = "OracleBot"
+        self.thread_id = 1 #TODO fix
+        self.config = create_agent_config(self.name, self.thread_id)
+        self.graph = self._build_agent(self.name)
 
-def get_agent():
-    """
-    Get our LangGraph agent with the given model and tools.
-    """
-   
-    class GraphConfig(TypedDict):
-        app_name: str
+    def answer_question(self, question: str):
+        """
+        Answer a question using the LangGraph agent.
+        """
+        messages = [HumanMessage(content=question)]
+        for mode, chunk in self.graph.stream({"messages": messages}, config=self.config, stream_mode=["messages", "custom"]): # type: ignore
+            print(chunk)
 
-    graph = StateGraph(State, context_schema=GraphConfig)
+    def _build_agent(self, name: str):
+        """
+        Get our LangGraph agent with the given model and tools.
+        """
     
-    # Add nodes
-    graph.add_node("agent", call_model)
-    graph.add_node("tools", tool_node)
-    
-    graph.add_edge(START, "agent")
-    graph.add_conditional_edges("agent", tools_condition)
-    graph.add_edge("tools", "agent")
-        
-    return graph.compile()
+        class GraphConfig(TypedDict):
+            name: str;
+            thread_id: int;
+
+        graph = StateGraph(state_schema=MessagesState, context_schema=GraphConfig)
+
+        # Add nodes
+        graph.add_node("agent", call_model)
+        graph.add_node("tools", tool_node)
+
+        # Add edges
+        graph.add_edge(START, "agent")
+        graph.add_conditional_edges("agent", tools_condition)
+        graph.add_edge("tools", "agent")
+
+        # Add memory
+        memory = InMemorySaver()
+            
+        return graph.compile(checkpointer=memory)
 
 # test
 if __name__ == "__main__":
     import os
     
     question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
-    messages = [HumanMessage(content=question)]
     
     try:
-        graph = get_agent()
-        from agent.config import create_agent_config
-        config = create_agent_config("OracleBot")
-        from app import start_phoenix
-        result = graph.invoke({"messages": messages}, config)
-        print("\n=== Agent Response ===")
-        for m in result["messages"]:
-            m.pretty_print()
+        from logger.phoenix import start_phoenix
+        start_phoenix()
+        bot = OracleBot()
+        bot.answer_question(question)
+
     except Exception as e:
         print(f"Error running agent: {e}")
         print("Make sure your API key is correct and the service is accessible.")

agent/config.py

@@ -1,6 +1,6 @@
 from typing import Literal
 
-def create_agent_config(app_name: str) -> dict:
+def create_agent_config(app_name: str, thread_id: int) -> dict:
     """
     Create configuration for the agent.
     
@@ -12,6 +12,7 @@ def create_agent_config(app_name: str) -> dict:
     """
     return {
         "configurable": {
-            "app_name": app_name
-        }
+            "app_name": app_name,
+            "thread_id": thread_id,
+        },
     }