Update agent.py

agent.py

@@ -1,45 +1,39 @@
-# agent.py
-from langchain_core.runnables import RunnableLambda, RunnableMap
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
+import operator
+import os
+import json # Required for json.loads in parse_llm_response
+from typing import TypedDict, Annotated, List, Dict, Any, Union
+from datetime import datetime
+
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate, SystemMessagePromptTemplate
+from langchain_core.runnables import RunnableLambda, RunnableMap # Added RunnableMap for input mapping
+
 from langgraph.graph import StateGraph, END
-import json
-from typing import TypedDict, List, Dict, Any, Union # Import TypedDict and other types for better type hints
 
-from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
-from tools import tools # Ensure 'tools' is a list of LangChain Tool objects
+# Use ChatGoogleGenerativeAI as per previous instructions for Gemini
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+# Import tools_for_llm from your tools.py file
+from tools import tools_for_llm # This imports the list of all your tools as tools_for_llm
+
+# Load environment variables
+from dotenv import load_dotenv
+load_dotenv()
+
+# Set up the LLM
+# Ensure your GOOGLE_API_KEY is set as an environment variable
+llm = ChatGoogleGenerativeAI(model="gemini-1.5-pro", temperature=0)
 
 # Load the system prompt from file
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    SYSTEM_PROMPT = f.read()
-
-# Define the shared memory structure for the agent
-# Use TypedDict for better type hinting and explicit state schema
-class AgentState(TypedDict):
-    """
-    Represents the state of the agent at each step of the graph.
-    All keys are required when creating an AgentState, but can be None or default values.
-    """
-    input: str # The original user query
-    chat_history: List[Union[HumanMessage, AIMessage, ToolMessage]] # History of messages
-    llm_response_raw: AIMessage # The raw AIMessage response from the LLM in the decision step
-    parsed_action: Dict[str, Any] # The dictionary with the parsed action (tool_name, tool_args, or final_answer)
-    tool_output: Any # The output from the execution of a tool (can be str or dict from tool)
-    output: str # The final answer from the agent to the user (formatted for GAIA)
-
-# Initialize the language model (using Hugging Face Inference Endpoint)
-# Ensure that this URL is correct and your HF token is available in the environment.
-endpoint = HuggingFaceEndpoint(
-    endpoint_url="https://api-inference.huggingface.co/models/Meta-Llama/llama-3-70b-instruct",
-    temperature=0.01, # A little temperature can aid creativity, but for tool-use decisions, 0 is common.
-    # It is crucial that your HF token is available as an environment variable HF_TOKEN
-)
-llm = ChatHuggingFace(llm=endpoint)
+    SYSTEM_PROMPT_CONTENT = f.read() # Renamed to avoid confusion with variable names
 
-# Helper to describe tools to the LLM
+# --- Helper to describe tools to the LLM ---
 def get_tool_descriptions(tool_list):
-    """Generates a string of tool descriptions for the LLM's prompt."""
+    """Generates a string of tool descriptions for the LLM's prompt.
+    This function generates descriptions without any special escaping for LangChain's templating,
+    as the entire string will be passed as a single variable.
+    """
     descriptions = []
     for tool_item in tool_list:
         args_schema_str = ""
@@ -64,20 +58,31 @@ def get_tool_descriptions(tool_list):
             except Exception:
                 args_schema_str = f"Arguments Schema: {tool_item.args_schema.__name__}"
 
-
         descriptions.append(f"- {tool_item.name}: {tool_item.description} {args_schema_str}")
     return "\n".join(descriptions)
 
-# Custom function to invoke tools by name and arguments
+# --- Agent State Definition ---
+class AgentState(TypedDict):
+    """
+    Represents the state of the agent at each step of the graph.
+    All keys are required when creating an AgentState, but can be None or default values.
+    """
+    input: str # The original user query
+    chat_history: Annotated[List[Union[HumanMessage, AIMessage, ToolMessage]], operator.add] # History of messages
+    llm_response_raw: Union[AIMessage, None] # The raw AIMessage response from the LLM in the decision step
+    parsed_action: Union[Dict[str, Any], None] # The dictionary with the parsed action (tool_name, tool_args, or final_answer)
+    tool_output: Union[Any, None] # The output from the execution of a tool (can be str or dict from tool)
+    output: Union[str, None] # The final answer from the agent to the user (formatted for GAIA)
+    tool_descriptions_str: str # Added to explicitly pass tool descriptions through state
+
+# --- Custom function to invoke tools by name and arguments ---
 def invoke_tool(tool_name: str, tool_args: dict):
     """Invokes a tool with its name and a dictionary of arguments."""
     print(f"[{__name__}] Invoking tool: '{tool_name}' with args: {tool_args}")
-    for tool_item in tools: # 'tools' here comes from the import tools.py
+    # Use tools_for_llm, which is imported from tools.py
+    for tool_item in tools_for_llm:
         if tool_item.name == tool_name:
             try:
-                # LangChain's @tool decorated functions are designed to work well
-                # when invoked with a dictionary of keyword arguments.
-                # The tool_item.invoke() method handles mapping these to function parameters.
                 result = tool_item.invoke(tool_args)
                 print(f"[{__name__}] Tool '{tool_name}' returned: {result}")
                 return result
@@ -85,95 +90,105 @@ def invoke_tool(tool_name: str, tool_args: dict):
                 return f"Error executing tool '{tool_name}' with args {tool_args}: {str(e)}"
     raise ValueError(f"Tool '{tool_name}' not found")
 
+# --- Nodes in the Agent Graph ---
+
 # 4. Node: Initial LLM call - Decide if tool is needed AND what arguments to use
 def call_llm_decide(state: AgentState) -> AgentState:
     """
     Prompts the LLM to decide on a tool and its arguments, or provide a direct answer.
     The LLM is expected to output a JSON string.
     """
-    print(f"[{__name__}] call_llm_decide: Initial state received (keys): {list(state.keys())}") # DEBUG PRINT
-    current_input = state["input"]
-    chat_history = state.get("chat_history", []) # Use .get for optional keys initially (though 'chat_history' should always be there from initial_state)
-
-    tool_descriptions = get_tool_descriptions(tools)
+    print(f"[{__name__}] call_llm_decide: Initial state received (keys): {list(state.keys())}")
 
-    # Inject tool_descriptions into the SYSTEM_PROMPT
-    # The SYSTEM_PROMPT now has a placeholder {{tool_descriptions}}
-    formatted_system_prompt = SYSTEM_PROMPT.replace("{{tool_descriptions}}", tool_descriptions)
+    current_input = state["input"]
+    chat_history = state.get("chat_history", [])
+    
+    # Generate the string of tool descriptions
+    tool_descriptions_string = get_tool_descriptions(tools_for_llm)
 
-    # CRUCIAL: The prompt needs to guide the LLM to output a specific JSON format
-    # Llama 3 Instruct is generally good at following this if trained or prompted well.
+    # Define the ChatPromptTemplate.
+    # 'tool_descriptions' is now an explicit input variable for the template.
     decision_prompt_template = ChatPromptTemplate.from_messages([
-        ("system", formatted_system_prompt), # Use the formatted system prompt here
+        SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT_CONTENT), # SYSTEM_PROMPT_CONTENT contains {{tool_descriptions}}
         MessagesPlaceholder(variable_name="chat_history"),
-        ("human", "{input}"),
+        HumanMessagePromptTemplate.from_template("{input}"),
     ])
 
-    chain = decision_prompt_template | llm
+    # Construct the chain for the LLM decision using RunnableMap for explicit input passing
+    chain = (
+        RunnableMap({
+            "input": operator.itemgetter("input"),
+            "chat_history": operator.itemgetter("chat_history"),
+            # Pass the generated tool descriptions string as the 'tool_descriptions' variable
+            "tool_descriptions": RunnableLambda(lambda x: tool_descriptions_string)
+        })
+        | decision_prompt_template
+        | llm
+    )
+
     response = chain.invoke({
         "input": current_input,
-        "chat_history": chat_history
+        "chat_history": chat_history,
+        # 'tool_descriptions' is handled by the RunnableMap above, no need to pass here directly
     })
 
     print(f"[{__name__}] LLM raw decision response: {response.content}")
 
-    # Parse the LLM's JSON output
-    parsed_action: Dict[str, Any] = {} # Explicitly type as Dict[str, Any]
+    parsed_action: Dict[str, Any] = {}
     try:
         parsed_action = json.loads(response.content)
     except json.JSONDecodeError as e:
         print(f"[{__name__}] Error parsing LLM JSON output: {e}. Raw content: {response.content}")
-        # Fallback: if parsing fails, assume it's a direct answer or an error.
-        # This can happen if the LLM doesn't follow the JSON format strictly.
         parsed_action = {"action": "final_answer", "answer": f"Error: Could not parse LLM's action. Raw LLM output: {response.content[:200]}..."}
 
-    # Add the LLM's raw response and parsed action to history for context
-    # Note: We add the HumanMessage here so the chat_history for the LLM is complete before the next step.
     new_state = AgentState(
-        input=current_input, # Always carry over input
+        input=current_input,
         chat_history=chat_history + [HumanMessage(content=current_input), response],
-        llm_response_raw=response, # Store the full AIMessage
-        parsed_action=parsed_action, # Store the parsed dictionary
-        tool_output="", # Initialize with empty string, will be populated by call_tool_node if needed
-        output="" # Initialize output, will be populated by generate_final_answer
+        llm_response_raw=response,
+        parsed_action=parsed_action,
+        tool_output=None, # Initialize with None
+        output=None, # Initialize with None
+        tool_descriptions_str=tool_descriptions_string # Store for subsequent nodes if needed
     )
-    print(f"[{__name__}] call_llm_decide: State being returned (keys): {list(new_state.keys())}") # DEBUG PRINT
+    print(f"[{__name__}] call_llm_decide: State being returned (keys): {list(new_state.keys())}")
     return new_state
 
 # 5. Node: Execute Tool (if chosen)
 def call_tool_node(state: AgentState) -> AgentState:
     """Executes the chosen tool with the provided arguments."""
-    print(f"[{__name__}] call_tool_node: State received (keys): {list(state.keys())}") # DEBUG PRINT
+    print(f"[{__name__}] call_tool_node: State received (keys): {list(state.keys())}")
+
     parsed_action = state["parsed_action"]
     tool_name = parsed_action.get("tool_name")
     tool_args = parsed_action.get("tool_args", {})
-
-    tool_output = ""
+    tool_output = None # Initialize as None
     tool_message = None
 
-    if tool_name: # If a tool was indicated in the parsed action
+    if tool_name:
         try:
-            tool_output = invoke_tool(tool_name, tool_args)
-            tool_message = ToolMessage(content=str(tool_output), name=tool_name)
+            raw_tool_output = invoke_tool(tool_name, tool_args)
+            # Ensure tool_output is always a string for ToolMessage content
+            tool_output = str(raw_tool_output)
+            tool_message = ToolMessage(content=tool_output, name=tool_name)
         except Exception as e:
-            tool_output = f"Error calling tool '{tool_name}': {str(e)}"
+            tool_output = f"Error executing tool '{tool_name}': {str(e)}"
             tool_message = ToolMessage(content=tool_output, name="tool_error")
     else:
-        # Fallback if somehow reached here without a tool_name
         tool_output = "No tool name provided by LLM, but reached tool execution node."
         tool_message = ToolMessage(content=tool_output, name="no_tool_fallback")
-
+    
     print(f"[{__name__}] Tool '{tool_name}' output: {tool_output}")
 
     new_state = AgentState(
-        input=state["input"], # Carry over from previous state
-        chat_history=state["chat_history"] + [tool_message], # Add the tool's output to history
-        llm_response_raw=state["llm_response_raw"], # Carry over from previous state
-        parsed_action=state["parsed_action"], # Carry over from previous state
-        tool_output=tool_output, # Store the actual output
-        output="" # Will be populated by generate_final_answer
+        input=state["input"],
+        chat_history=state["chat_history"] + [tool_message],
+        llm_response_raw=state["llm_response_raw"],
+        parsed_action=state["parsed_action"],
+        tool_output=tool_output,
+        output=None, # Will be populated by generate_final_answer
+        tool_descriptions_str=state["tool_descriptions_str"]
     )
-    print(f"[{__name__}] call_tool_node: State being returned (keys): {list(new_state.keys())}") # DEBUG PRINT
+    print(f"[{__name__}] call_tool_node: State being returned (keys): {list(new_state.keys())}")
     return new_state
 
 # 6. Node: Generate Final Answer (after tool or direct)
@@ -182,86 +197,90 @@ def generate_final_answer(state: AgentState) -> AgentState:
     Generates the final answer based on the conversation history,
     including any tool outputs, and formats it for GAIA evaluation.
     """
-    print(f"[{__name__}] generate_final_answer: State received (keys): {list(state.keys())}") # DEBUG PRINT
+    print(f"[{__name__}] generate_final_answer: State received (keys): {list(state.keys())}")
     current_input = state["input"]
-    chat_history = state["chat_history"] # chat_history should already be populated from previous nodes
-    parsed_action = state["parsed_action"] # Get the action from the initial decision
+    chat_history = state["chat_history"]
+    parsed_action = state["parsed_action"]
+    tool_descriptions_string = state["tool_descriptions_str"] # Use the stored descriptions
 
     final_answer_content = ""
 
-    # If the LLM already gave a final answer directly in the initial decision, use that.
-    if parsed_action.get("action") == "final_answer" and "answer" in parsed_action:
+    if parsed_action and parsed_action.get("action") == "final_answer" and "answer" in parsed_action:
         final_answer_content = parsed_action["answer"]
         print(f"[{__name__}] Using direct final answer from initial LLM decision: {final_answer_content}")
     else:
         # If a tool was used, or if the LLM didn't give a direct answer initially,
         # prompt the LLM again to synthesize the final answer.
-        # The chat_history now contains the tool call and its output.
-        # This prompt should guide the LLM to give the answer content ONLY.
         final_prompt_template = ChatPromptTemplate.from_messages([
-            ("system", SYSTEM_PROMPT.replace("{{tool_descriptions}}", get_tool_descriptions(tools)) + "\n\n" # Re-inject descriptions
-                                       "You have processed the user's request, possibly using tools. "
-                                       "Now, synthesize all available information, including the results of any tool executions, "
-                                       "to provide a concise and direct final answer to the user. "
-                                       "If a tool was executed, use its output to formulate the answer. "
-                                       "If no tool was executed (or it failed), try to answer based on your general knowledge or indicate the issue."
-                                       "Crucially, provide ONLY the answer content, without any prefix, conversational filler, or apologies."
-                                       "For numbers, do not include commas, currency symbols, or units unless explicitly requested."
-                                       "For strings, write digits in plain text unless otherwise specified."
-                                       "The final output will be prefixed by 'FINAL ANSWER:' by the system."),
+            SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT_CONTENT), # SYSTEM_PROMPT_CONTENT contains {{tool_descriptions}}
             MessagesPlaceholder(variable_name="chat_history"), # This history now includes the tool call and its output
-            ("human", "{input}"), # Still pass the original input for context
-            # Explicitly indicate that the LLM should give only the answer
-            AIMessage(content="Based on the above, your concise answer is:") # This acts as a strong steer for LLM to provide only the answer
+            HumanMessagePromptTemplate.from_template("{input}"), # Still pass the original input for context
+            AIMessage(content="Based on the above, your concise answer is:") # Strong steer for LLM
         ])
-        chain = final_prompt_template | llm | StrOutputParser()
+        
+        # Build chain for final answer generation, including RunnableMap
+        final_answer_chain = (
+            RunnableMap({
+                "input": operator.itemgetter("input"),
+                "chat_history": operator.itemgetter("chat_history"),
+                "tool_descriptions": RunnableLambda(lambda x: tool_descriptions_string) # Pass the generated descriptions
+            })
+            | final_prompt_template
+            | llm
+            | StrOutputParser()
+        )
+
         print(f"[{__name__}] Generating final answer content with chat history and tool output...")
-        final_answer_content = chain.invoke({
+        final_answer_content = final_answer_chain.invoke({
             "input": current_input,
             "chat_history": chat_history # Full history including tool outputs
         })
         print(f"[{__name__}] Generated raw final answer content: {final_answer_content}")
 
-    # **CRUCIAL GAIA REQUIREMENT: Add the "FINAL ANSWER:" prefix here**
+    # CRUCIAL GAIA REQUIREMENT: Add the "FINAL ANSWER:" prefix here
     gaia_formatted_answer = f"FINAL ANSWER: {final_answer_content.strip()}"
     print(f"[{__name__}] GAIA formatted final answer: {gaia_formatted_answer}")
 
     new_state = AgentState(
-        input=current_input, # Carry over from previous state
+        input=current_input,
         chat_history=chat_history + [AIMessage(content=gaia_formatted_answer)],
-        llm_response_raw=state.get("llm_response_raw", AIMessage(content="")), # Use .get with default for robustness
-        parsed_action=state.get("parsed_action", {}), # Use .get with default
-        tool_output=state.get("tool_output", ""), # Use .get with default
-        output=gaia_formatted_answer # This is the string returned to the user/GAIA
+        llm_response_raw=state.get("llm_response_raw", AIMessage(content="")),
+        parsed_action=state.get("parsed_action", {}),
+        tool_output=state.get("tool_output", None), # Use None as default
+        output=gaia_formatted_answer,
+        tool_descriptions_str=tool_descriptions_string
     )
-    print(f"[{__name__}] generate_final_answer: State being returned (keys): {list(new_state.keys())}") # DEBUG PRINT
+    print(f"[{__name__}] generate_final_answer: State being returned (keys): {list(new_state.keys())}")
     return new_state
 
 # 7. Router: Decide which path to take after initial LLM call
 def route_action(state: AgentState) -> str:
     """Routes the graph based on the LLM's parsed action."""
-    print(f"[{__name__}] route_action: State received (keys): {list(state.keys())}") # DEBUG PRINT
-    parsed_action = state["parsed_action"]
-    action_type = parsed_action.get("action")
+    print(f"[{__name__}] route_action: State received (keys): {list(state.keys())}")
+    parsed_action = state.get("parsed_action") # Use .get for robustness
+
+    if not parsed_action: # Handle cases where parsing might have failed
+        print(f"[{__name__}] Parsed action is empty or None. Defaulting to 'generate_final_answer'.")
+        return "generate_final_answer" # Route to final answer, likely an error explanation
 
+    action_type = parsed_action.get("action")
     if action_type == "tool":
         print(f"[{__name__}] Routing to 'execute_tool' based on LLM decision.")
         return "execute_tool"
     elif action_type == "final_answer":
         print(f"[{__name__}] Routing to 'generate_final_answer' directly based on LLM decision.")
-        return "generate_final_answer" # Direct to final answer node if LLM decided directly
+        return "generate_final_answer"
     else:
-        print(f"[{__name__}] Unknown action type '{action_type}' or no action. Defaulting to 'generate_final_answer'.")
-        # Fallback if the LLM didn't produce a recognizable action
+        print(f"[{__name__}] Unknown action type '{action_type}'. Defaulting to 'generate_final_answer'.")
         return "generate_final_answer" # Route to final answer, might be an error case
 
 # 8. Build the agent graph
-builder = StateGraph(AgentState) # Use the updated AgentState
+builder = StateGraph(AgentState)
 
 # Add nodes
 builder.add_node("initial_llm_decision", call_llm_decide)
 builder.add_node("execute_tool", call_tool_node)
-builder.add_node("generate_final_answer", generate_final_answer) # This node generates final answer
+builder.add_node("generate_final_answer", generate_final_answer)
 
 # Set the entry point
 builder.set_entry_point("initial_llm_decision")
@@ -288,21 +307,23 @@ agent_executor = builder.compile()
 class BasicAgent:
     def __init__(self):
         self.agent = agent_executor
-
+    
     def __call__(self, question: str) -> str:
         # The initial state for each new question
         initial_state: AgentState = { # Explicitly type the dictionary as AgentState
             "input": question,
             "chat_history": [],
-            # Provide default empty/dummy values for all TypedDict keys
-            "llm_response_raw": AIMessage(content=""),
-            "parsed_action": {},
-            "tool_output": "",
-            "output": ""
+            "llm_response_raw": None, # Should be initialized as None
+            "parsed_action": None, # Should be initialized as None
+            "tool_output": None, # Should be initialized as None
+            "output": None, # Should be initialized as None
+            "tool_descriptions_str": "" # Initialize this new key
         }
+        
         # The invoke will run through the graph
         # LangGraph's invoke returns the final state
-        final_state = self.agent.invoke(initial_state) # Pass the explicitly typed initial state
+        final_state = self.agent.invoke(initial_state)
+        
         # Extract the final output from the state
         return final_state.get("output", "No answer could be generated.")
 
@@ -310,35 +331,30 @@ if __name__ == "__main__":
     # Example Usage (for local testing)
     print("Testing BasicAgent locally...")
 
-    # For local testing, ensure you have:
-    # 1. A mock or real `tools.py` that provides `tools` (a list of LangChain Tool objects).
-    # 2. A `system_prompt.txt` file (create it if not present, with example content below).
-    # 3. Your Hugging Face token set as an environment variable (HF_TOKEN).
-
     try:
         agent = BasicAgent()
+        
         print("\n--- Test 1: Simple question, should directly answer ---")
         response1 = agent("What is the capital of France?")
         print(f"Agent Response: {response1}")
-
+        
         print("\n--- Test 2: Question requiring a tool (e.g., web_search) ---")
-        # This will test if the agent can correctly call web_search
         # Make sure TAVILY_API_KEY is set in your environment for this to work
         response2 = agent("What is the current population of the United States? (as of today)")
         print(f"Agent Response: {response2}")
-
-        print("\n--- Test 3: Math question (e.g., multiply) ---")
+        
+        print("\n--- Test 3: Math question (e.g., calculator tool) ---")
         response3 = agent("What is 15 multiplied by 23?")
         print(f"Agent Response: {response3}")
-
+        
         print("\n--- Test 4: Question requiring Python code execution ---")
         response4 = agent("What is the result of the python code: `sum([x**2 for x in range(1, 5)])`?")
         print(f"Agent Response: {response4}")
-
+        
         print("\n--- Test 5: Question with no clear tool, but needs a general answer ---")
         response5 = agent("What is the meaning of life?")
         print(f"Agent Response: {response5}")
 
     except Exception as e:
         print(f"\nError during local testing: {e}")
-        print("Please ensure your HF_TOKEN and TAVILY_API_KEY are set, and 'tools.py' is correctly implemented.")
+        print("Please ensure your GOOGLE_API_KEY and TAVILY_API_KEY are set, and 'tools.py' is correctly implemented and exports 'tools_for_llm'.")