Update agent.py

agent.py

@@ -1,253 +1,146 @@
+# ==============================================================================
+# 1. IMPORTS AND SETUP
+# ==============================================================================
 import os
 from dotenv import load_dotenv
-import gradio as gr
+from typing import TypedDict, Annotated, List
 
-
-# LangGraph & LangChain
-from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import ToolNode, tools_condition
-from langchain_core.messages import SystemMessage, HumanMessage
-from langchain_core.tools import tool
-
-#infrence provider
+# LangChain and LangGraph imports
 from langchain_huggingface import HuggingFaceEndpoint
-# Web search tool
 from langchain_community.tools.tavily_search import TavilySearchResults
-
-# # NEW IMPORT
-# from langchain_experimental.tools import PythonREPLTool
-# from langchain_core.messages import BaseMessage, HumanMessage
-# from langgraph.graph import StateGraph, END
-# from langgraph.prebuilt import ToolNode
-# from typing import TypedDict, Annotated, List
-
-
-# --- 1. LOAD API KEYS ---
+from langchain_experimental.tools import PythonREPLTool
+from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolNode
+
+# ==============================================================================
+# 2. LOAD API KEYS AND DEFINE TOOLS
+# ==============================================================================
 load_dotenv()
 hf_token = os.getenv("HF_TOKEN")
 tavily_api_key = os.getenv("TAVILY_API_KEY")
 
 if not hf_token or not tavily_api_key:
-    raise ValueError("Hugging Face Token or Tavily API Key is not set in the environment variables.")
+    # This will show a clear error in the logs if keys are missing
+    raise ValueError("HF_TOKEN or TAVILY_API_KEY not set. Please add them to your Space secrets.")
 os.environ["TAVILY_API_KEY"] = tavily_api_key
 
-
-# --- 2. DEFINE TOOLS and INITIALIZE LLM ---
-# UPDATED TOOLS LIST
-# tools = [TavilySearchResults(max_results=3), PythonREPLTool()]
-# tool_node = ToolNode(tools)
-
-
-### TOOLS
-
-@tool
-def multiply(a: int, b: int) -> int:
-    """Multiply two numbers.
-    Args:
-        a: first int
-        b: second int
-    """
-    return a * b
-
-@tool
-def add(a: int, b: int) -> int:
-    """Add two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
-    return a + b
-
-@tool
-def subtract(a: int, b: int) -> int:
-    """Subtract two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
-    return a - b
-
-@tool
-def divide(a: int, b: int) -> int:
-    """Divide two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
-    if b == 0:
-        raise ValueError("Cannot divide by zero.")
-    return a / b
-
-@tool
-def modulus(a: int, b: int) -> int:
-    """Get the modulus of two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
-    return a % b
-
-
-
-@tool
-def web_search(query: str) -> str:
-    """Search Tavily for a query and return maximum 3 results.
-    
-    Args:
-        query: The search query."""
-    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"web_results": formatted_search_docs}
-
-# SYSTEM PROMPT
-system_prompt = """
-You are a helpful assistant tasked with answering questions using a set of tools. 
-Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
-FINAL ANSWER: [YOUR FINAL ANSWER]. 
-YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
-Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
+# The agent's tools
+tools = [TavilySearchResults(max_results=3, description="A search engine for finding up-to-date information on the web."), PythonREPLTool()]
+tool_node = ToolNode(tools)
+
+# ==============================================================================
+# 3. CONFIGURE THE LLM (THE "BRAIN")
+# ==============================================================================
+# The model we'll use as the agent's brain
+repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+
+# The system prompt gives the agent its mission and instructions
+SYSTEM_PROMPT = """You are a highly capable AI agent named 'GAIA-Solver'. Your mission is to accurately answer complex questions.
+
+**Your Instructions:**
+1.  **Analyze:** Carefully read the user's question to understand all parts of what is being asked.
+2.  **Plan:** Think step-by-step. Break the problem into smaller tasks. Decide which tool is best for each task. (e.g., use 'tavily_search_results_json' for web searches, use 'python_repl' for calculations or code execution).
+3.  **Execute:** Call ONE tool at a time.
+4.  **Observe & Reason:** After getting a tool's result, observe it. Decide if you have the final answer or if you need to use another tool.
+5.  **Final Answer:** Once you are confident, provide a clear, direct, and concise final answer. Do not include your thought process in the final answer.
 """
-tools = [divide, add, multiply,subtract, web_search,modulus ]
-
-### LLM
-# repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-# llm = HuggingFaceEndpoint(
-#     repo_id=repo_id,
-#     huggingfacehub_api_token=hf_token,
-#     temperature=0.1,
-#     max_new_tokens=1024,
-# )
-# llm_with_tools = llm.bind_tools(tools)
-
-
-def build_graph():
-    """Builds and returns the LangGraph graph."""
-    #llm = ChatGroq(model="qwen-qwq-32b", temperature=0,api_key=groq_api_key)
-    repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-    llm = HuggingFaceEndpoint(
+
+# Initialize the LLM endpoint
+llm = HuggingFaceEndpoint(
     repo_id=repo_id,
     huggingfacehub_api_token=hf_token,
-    temperature=0,
-    max_new_tokens=1024,
+    temperature=0, # Set to 0 for deterministic, less random output
+    max_new_tokens=2048,
 )
-    llm_with_tools = llm.bind_tools(tools)
-
-# Node
-    def assistant(state: MessagesState):
-        """Assistant node"""
-        return {"messages": [llm_with_tools.invoke([system_prompt] + state["messages"])]}
-
-
-    builder = StateGraph(MessagesState)
-    # Nodes
-    builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
-
-    # Edges
-    builder.add_edge(START, "assistant")
-    builder.add_conditional_edges("assistant", tools_condition)
-    builder.add_edge("tools", "assistant")
-
-    #Compile graph
-    return builder.compile()
-
-
-
-
-
 
+# ==============================================================================
+# 4. BUILD THE LANGGRAPH AGENT
+# ==============================================================================
 
-
-
-
-
-
-
-
-
-
-
-# --- 3. DEFINE THE AGENT'S STATE ---
-"""
+# Define the Agent's State (its memory)
 class AgentState(TypedDict):
     messages: Annotated[List[BaseMessage], lambda x, y: x + y]
 
+# This is a more robust way to combine the prompt, model, and tool binding
+# It ensures the system prompt is always used.
+llm_with_tools = llm.bind_tools(tools)
 
+# Define the Agent Node
+def agent_node(state):
+    # Get the last message to pass to the model
+    last_message = state['messages'][-1]
 
+    # Prepend the system prompt to every call
+    prompt_with_system = [
+        HumanMessage(content=SYSTEM_PROMPT, name="system_prompt"),
+        last_message
+    ]
 
-
-# --- 4. DEFINE THE NODES OF THE GRAPH ---
-def agent_node(state):
-    response = llm_with_tools.invoke(state["messages"])
+    response = llm_with_tools.invoke(prompt_with_system)
     return {"messages": [response]}
 
-
-# --- 5. DEFINE THE EDGES OF THE GRAPH ---
+# Define the Edge Logic
 def should_continue(state):
     last_message = state["messages"][-1]
     if last_message.tool_calls:
-        return "tools"
-    return END
-
+        return "tools" # Route to the tool node
+    return END # End the process
 
-# --- 6. ASSEMBLE THE GRAPH ---
+# Assemble the graph
 workflow = StateGraph(AgentState)
-
 workflow.add_node("agent", agent_node)
 workflow.add_node("tools", tool_node)
-
 workflow.set_entry_point("agent")
-
 workflow.add_conditional_edges(
     "agent",
     should_continue,
-    {
-        "tools": "tools",
-        "end": END,
-    },
+    {"tools": "tools", "end": END},
 )
-
 workflow.add_edge("tools", "agent")
 
+# Compile the graph into a runnable app
 app = workflow.compile()
 
 
-# --- 7. CREATE THE USER INTERFACE (UI) ---
-def run_agent(query: str):
-    try:
-        inputs = {"messages": [HumanMessage(content=query)]}
-        final_response = None
-        # Using stream to get final output, can be slow for complex tasks
-        for s in app.stream(inputs, {"recursion_limit": 10}):
-            if "agent" in s:
-                final_response = s["agent"]["messages"][-1].content
-        return final_response if final_response else "Agent did not produce a final answer."
-
-    except Exception as e:
-        return f"An error occurred: {e}"
-
-iface = gr.Interface(
-    fn=run_agent,
-    inputs=gr.Textbox(lines=2, placeholder="Ask the agent anything..."),
-    outputs="markdown",
-    title="GAIA Agent v0.3 (LangGraph + Code Interpreter)",
-    description="This agent can use web search and a Python code interpreter.",
-    examples=[
-        ["What is the square root of the number of states in the USA?"],
-        ["What is the total number of letters in the names of the first three planets in our solar system?"]
-    ],
-)
-
-
-# --- 8. LAUNCH THE APP ---
-iface.launch()
-# gr.launch()
-"""
+# ==============================================================================
+# 5. THE BASICAGENT CLASS (FOR THE TEST HARNESS)
+# This MUST be at the end, after `app` is defined.
+# ==============================================================================
+class BasicAgent:
+    """
+    This is the agent class that the GAIA test harness will use.
+    """
+    def __init__(self):
+        # The compiled LangGraph app is our agent executor
+        self.agent_executor = app
+
+    def run(self, question: str) -> str:
+        """
+        This method is called by the test script with each question.
+        It runs the LangGraph agent and returns the final answer.
+        """
+        print(f"Agent received question (first 80 chars): {question[:80]}...")
+        try:
+            # Format the input for our graph
+            inputs = {"messages": [HumanMessage(content=question)]}
+
+            # Stream the response to get the final answer
+            final_response = ""
+            for s in self.agent_executor.stream(inputs, {"recursion_limit": 15}):
+                if "agent" in s:
+                    # The final answer is the content of the last message from the agent node
+                    if s["agent"]["messages"][-1].content:
+                         final_response = s["agent"]["messages"][-1].content
+
+            # A fallback in case the agent finishes without a clear message
+            if not final_response:
+                final_response = "Agent finished but did not produce a final answer."
+
+            print(f"Agent returning final answer (first 80 chars): {final_response[:80]}...")
+            return final_response
+
+        except Exception as e:
+            print(f"An error occurred in agent execution: {e}")
+            return f"Error: {e}"