Update agent.py

agent.py

@@ -36,37 +36,13 @@ 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 = build_graph(self.tools, self.agent_executor)  # Using the standalone function
-
+        self.workflow = self._build_workflow()  # Correct attribute name
+    
     def _build_agent_executor(self) -> AgentExecutor:
         """Build the agent executor"""
         prompt = ChatPromptTemplate.from_messages([
@@ -135,17 +111,4 @@ class AIAgent:
             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
-
-if __name__ == "__main__":
-    agent = AIAgent()
-    response = agent("What's the capital of France?")
-    print("Response:", response["response"])
-    if response["sources"]:
-        print("\nSources:")
-        for source in response["sources"]:
-            print("-", source)
-    if response["steps"]:
-        print("\nSteps:")
-        for step in response["steps"]:
-            print("-", step)
+        return steps