Update agent.py

agent.py

@@ -1,7 +1,7 @@
 # agent.py
 import os
 from dotenv import load_dotenv
-from typing import TypedDict, Annotated, Sequence
+from typing import TypedDict, Annotated, Sequence, Dict, Any
 from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
@@ -9,8 +9,7 @@ from langgraph.graph import END, StateGraph
 from langgraph.prebuilt import ToolNode
 from langchain_community.tools import DuckDuckGoSearchResults
 from langchain_community.utilities import WikipediaAPIWrapper
-from langchain.agents import create_tool_calling_agent
-from langchain.agents import AgentExecutor
+from langchain.agents import create_tool_calling_agent, AgentExecutor
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 import operator
 
@@ -37,43 +36,53 @@ def calculate(expression: str) -> str:
     python_repl = PythonREPL()
     return python_repl.run(expression)
 
-def build_agent(tools: list, llm: ChatOpenAI) -> AgentExecutor:
-    """Build the agent executor"""
-    prompt = ChatPromptTemplate.from_messages([
-        ("system", "You are a helpful AI assistant. Use tools when needed."),
-        MessagesPlaceholder(variable_name="messages"),
-        MessagesPlaceholder(variable_name="agent_scratchpad"),
-    ])
-    agent = create_tool_calling_agent(llm, tools, prompt)
-    return AgentExecutor(agent=agent, tools=tools, verbose=True)
-
-def build_graph(tools: list, agent: AgentExecutor) -> StateGraph:
-    """Build the LangGraph workflow"""
-    workflow = StateGraph(AgentState)
+class AIAgent:
+    def __init__(self, model_name: str = "gpt-3.5-turbo"):
+        self.tools = [wikipedia_search, web_search, calculate]
+        self.llm = ChatOpenAI(model=model_name, temperature=0.7)
+        self.agent_executor = self._build_agent_executor()
+        self.workflow = self._build_workflow()  # Changed from 'graph' to 'workflow'
     
-    # Define nodes
-    workflow.add_node("agent", lambda state: {"messages": [agent.invoke(state)["output"]]})
-    workflow.add_node("tools", ToolNode(tools))
+    def _build_agent_executor(self) -> AgentExecutor:
+        """Build the agent executor"""
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", "You are a helpful AI assistant. Use tools when needed."),
+            MessagesPlaceholder(variable_name="messages"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ])
+        agent = create_tool_calling_agent(self.llm, self.tools, prompt)
+        return AgentExecutor(agent=agent, tools=self.tools, verbose=True)
     
-    # Define edges
-    workflow.set_entry_point("agent")
-    workflow.add_conditional_edges(
-        "agent",
-        lambda state: "continue" if state["messages"][-1].additional_kwargs.get("tool_calls") else "end",
-        {"continue": "tools", "end": END}
-    )
-    workflow.add_edge("tools", "agent")
+    def _build_workflow(self) -> StateGraph:
+        """Build and return the compiled workflow"""
+        workflow = StateGraph(AgentState)
+        
+        # Define nodes
+        workflow.add_node("agent", self._run_agent)
+        workflow.add_node("tools", ToolNode(self.tools))
+        
+        # Define edges
+        workflow.set_entry_point("agent")
+        workflow.add_conditional_edges(
+            "agent",
+            self._should_continue,
+            {"continue": "tools", "end": END}
+        )
+        workflow.add_edge("tools", "agent")
+        
+        return workflow.compile()
     
-    return workflow.compile()
-
-class AdvancedAIAgent:
-    def __init__(self, model_name: str = "gpt-4-turbo"):
-        self.tools = [wikipedia_search, web_search, calculate]
-        self.llm = ChatOpenAI(model=model_name, temperature=0.7)
-        self.agent = build_agent(self.tools, self.llm)
-        self.workflow = build_graph(self.tools, self.agent)
+    def _run_agent(self, state: AgentState) -> Dict[str, Any]:
+        """Execute the agent"""
+        response = self.agent_executor.invoke({"messages": state["messages"]})
+        return {"messages": [response["output"]]}
     
-    def __call__(self, query: str) -> dict:
+    def _should_continue(self, state: AgentState) -> str:
+        """Determine if the workflow should continue"""
+        last_message = state["messages"][-1]
+        return "continue" if last_message.additional_kwargs.get("tool_calls") else "end"
+    
+    def __call__(self, query: str) -> Dict[str, Any]:
         """Process a user query"""
         state = AgentState(messages=[HumanMessage(content=query)], sender="user")
         
@@ -89,23 +98,19 @@ class AdvancedAIAgent:
                             }
         return {"response": "No response generated", "sources": [], "steps": []}
     
-    def _extract_sources(self, messages: Sequence[BaseMessage]) -> list:
+    def _extract_sources(self, messages: Sequence[BaseMessage]) -> List[str]:
+        """Extract sources from tool messages"""
         return [
             f"{msg.additional_kwargs.get('name', 'unknown')}: {msg.content}"
             for msg in messages
             if hasattr(msg, 'additional_kwargs') and 'name' in msg.additional_kwargs
         ]
     
-    def _extract_steps(self, messages: Sequence[BaseMessage]) -> list:
+    def _extract_steps(self, messages: Sequence[BaseMessage]) -> List[str]:
+        """Extract reasoning steps"""
         steps = []
         for msg in messages:
             if hasattr(msg, 'additional_kwargs') and 'tool_calls' in msg.additional_kwargs:
                 for call in msg.additional_kwargs['tool_calls']:
                     steps.append(f"Used {call['function']['name']}: {call['function']['arguments']}")
-        return steps
-
-# Example usage
-if __name__ == "__main__":
-    agent = AdvancedAIAgent()
-    response = agent("What's the capital of France?")
-    print(response)
+        return steps