Update agent.py

agent.py

@@ -1,3 +1,4 @@
+# agent.py
 from typing import TypedDict, Annotated, Sequence
 from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_core.tools import tool
@@ -10,14 +11,11 @@ from langchain.agents import create_tool_calling_agent
 from langchain.agents import AgentExecutor
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 import operator
-import json
 
-# Define the agent state
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], operator.add]
     sender: str
 
-# Initialize tools
 @tool
 def wikipedia_search(query: str) -> str:
     """Search Wikipedia for information."""
@@ -31,75 +29,50 @@ def web_search(query: str, num_results: int = 3) -> list:
 @tool
 def calculate(expression: str) -> str:
     """Evaluate mathematical expressions."""
-    from langchain_experimental.tools import PythonREPLTool
-    python_repl = PythonREPLTool()
-    return python_repl.run(f"print({expression})")
+    from langchain_experimental.utilities import PythonREPL
+    python_repl = PythonREPL()
+    return python_repl.run(expression)
 
-class AdvancedAIAgent:
-    def __init__(self, model_name="gpt-4-turbo"):
-        # Initialize tools and LLM
-        self.tools = [wikipedia_search, web_search, calculate]
-        self.llm = ChatOpenAI(model=model_name, temperature=0.7)
-        
-        # Create the agent
-        self.agent = self._create_agent()
-        
-        # Build the graph workflow
-        self.workflow = self._build_graph()
+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 _create_agent(self) -> AgentExecutor:
-        """Create the agent with tools and prompt"""
-        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)
-
-    def _build_graph(self):
-        """Build the LangGraph workflow"""
-        workflow = StateGraph(AgentState)
-        
-        # Define nodes
-        workflow.add_node("agent", self._call_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()
-
-    def _call_agent(self, state: AgentState):
-        """Execute the agent"""
-        response = self.agent.invoke({"messages": state["messages"]})
-        return {"messages": [response["output"]]}
-
-    def _should_continue(self, state: AgentState):
-        """Determine if the workflow should continue"""
-        last_message = state["messages"][-1]
-        
-        # If no tool calls, end
-        if not last_message.additional_kwargs.get("tool_calls"):
-            return "end"
-        return "continue"
+def build_graph(tools: list, agent: AgentExecutor) -> StateGraph:
+    """Build the LangGraph workflow"""
+    workflow = StateGraph(AgentState)
+    
+    # Define nodes
+    workflow.add_node("agent", lambda state: {"messages": [agent.invoke(state)["output"]]})
+    workflow.add_node("tools", ToolNode(tools))
+    
+    # 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")
+    
+    return workflow.compile()
 
+class AIAgent:
+    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 __call__(self, query: str) -> dict:
         """Process a user query"""
-        # Initialize state
         state = AgentState(messages=[HumanMessage(content=query)], sender="user")
         
-        # Execute the workflow
         for output in self.workflow.stream(state):
             for key, value in output.items():
                 if key == "messages":
@@ -110,20 +83,25 @@ class AdvancedAIAgent:
                                 "sources": self._extract_sources(state["messages"]),
                                 "steps": self._extract_steps(state["messages"])
                             }
-
+        return {"response": "No response generated", "sources": [], "steps": []}
+    
     def _extract_sources(self, messages: Sequence[BaseMessage]) -> list:
-        """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:
-        """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
+        return steps
+
+# Example usage
+if __name__ == "__main__":
+    agent = AIAgent()
+    response = agent("What's the capital of France?")
+    print(response)