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
+from typing import TypedDict, Annotated, Sequence, Dict, Any, List  # Added List here
 from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
@@ -12,6 +12,7 @@ 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()
 
@@ -41,7 +42,7 @@ class AIAgent:
         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'
+        self.workflow = self._build_workflow()
     
     def _build_agent_executor(self) -> AgentExecutor:
         """Build the agent executor"""
@@ -57,11 +58,9 @@ class AIAgent:
         """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",
@@ -113,4 +112,17 @@ 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
+        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)