Update agent.py

agent.py

@@ -3,15 +3,16 @@ from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
 from langgraph.graph import END, StateGraph
-from langgraph.prebuilt import ToolExecutor, ToolInvocation
+from langgraph.prebuilt import ToolNode
 from langchain.tools import DuckDuckGoSearchResults
 from langchain_community.utilities import WikipediaAPIWrapper
 from langchain.agents import create_tool_calling_agent
 from langchain.agents import AgentExecutor
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 import operator
+import json
 
-load_dotenv()
+load dotenv()
 
 # Define the agent state
 class AgentState(TypedDict):
@@ -44,7 +45,6 @@ class AdvancedAIAgent:
         
         # Create the agent
         self.agent = self._create_agent()
-        self.tool_executor = ToolExecutor(self.tools)
         
         # Build the graph workflow
         self.workflow = self._build_graph()
@@ -66,7 +66,7 @@ class AdvancedAIAgent:
         
         # Define nodes
         workflow.add_node("agent", self._call_agent)
-        workflow.add_node("tools", self._call_tools)
+        workflow.add_node("tools", ToolNode(self.tools))  # Using ToolNode instead of ToolExecutor
         
         # Define edges
         workflow.set_entry_point("agent")
@@ -87,31 +87,6 @@ class AdvancedAIAgent:
         response = self.agent.invoke({"messages": state["messages"]})
         return {"messages": [response["output"]]}
 
-    def _call_tools(self, state: AgentState):
-        """Execute tools"""
-        last_message = state["messages"][-1]
-        
-        # Find the tool calls
-        tool_calls = last_message.additional_kwargs.get("tool_calls", [])
-        
-        # Execute each tool
-        for tool_call in tool_calls:
-            action = ToolInvocation(
-                tool=tool_call["function"]["name"],
-                tool_input=json.loads(tool_call["function"]["arguments"]),
-            )
-            output = self.tool_executor.invoke(action)
-            
-            # Create tool message
-            tool_message = ToolMessage(
-                content=str(output),
-                name=action.tool,
-                tool_call_id=tool_call["id"],
-            )
-            state["messages"].append(tool_message)
-        
-        return {"messages": state["messages"]}
-
     def _should_continue(self, state: AgentState):
         """Determine if the workflow should continue"""
         last_message = state["messages"][-1]
@@ -153,27 +128,4 @@ class AdvancedAIAgent:
             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()
-    
-    queries = [
-        "What is the capital of France?",
-        "Calculate 15% of 200",
-        "Tell me about the latest developments in quantum computing"
-    ]
-    
-    for query in queries:
-        print(f"\nQuestion: {query}")
-        response = agent(query)
-        print(f"Answer: {response['response']}")
-        if response['sources']:
-            print("Sources:")
-            for source in response['sources']:
-                print(f"- {source}")
-        if response['steps']:
-            print("Steps taken:")
-            for step in response['steps']:
-                print(f"- {step}")
+        return steps