Update agent.py

agent.py

@@ -1,7 +1,7 @@
 # agent.py
 import os
 from dotenv import load_dotenv
-from typing import TypedDict, Annotated, Sequence, Dict, Any, List  # Added List here
+from typing import TypedDict, Annotated, Sequence, Dict, Any, List
 from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
@@ -12,7 +12,6 @@ from langchain_community.utilities import WikipediaAPIWrapper
 from langchain.agents import create_tool_calling_agent, AgentExecutor
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 import operator
-import operator
 
 load_dotenv()
 
@@ -37,13 +36,37 @@ def calculate(expression: str) -> str:
     python_repl = PythonREPL()
     return python_repl.run(expression)
 
+def build_graph(tools: list, agent_executor: AgentExecutor) -> StateGraph:
+    """Build and return the compiled workflow graph"""
+    workflow = StateGraph(AgentState)
+    
+    def run_agent(state: AgentState) -> Dict[str, Any]:
+        response = agent_executor.invoke({"messages": state["messages"]})
+        return {"messages": [response["output"]]}
+    
+    def should_continue(state: AgentState) -> str:
+        last_message = state["messages"][-1]
+        return "continue" if last_message.additional_kwargs.get("tool_calls") else "end"
+    
+    workflow.add_node("agent", run_agent)
+    workflow.add_node("tools", ToolNode(tools))
+    workflow.set_entry_point("agent")
+    workflow.add_conditional_edges(
+        "agent",
+        should_continue,
+        {"continue": "tools", "end": END}
+    )
+    workflow.add_edge("tools", "agent")
+    
+    return workflow.compile()
+
 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()
-    
+        self.workflow = build_graph(self.tools, self.agent_executor)  # Using the standalone function
+
     def _build_agent_executor(self) -> AgentExecutor:
         """Build the agent executor"""
         prompt = ChatPromptTemplate.from_messages([