Update app.py

app.py

@@ -1,157 +1,180 @@
 import os
 import gradio as gr
 import requests
+import inspect
 import pandas as pd
-import json
-import traceback
 from dotenv import load_dotenv
-from typing import List, Dict, Any, Optional
+from typing import TypedDict, Annotated, Sequence, List, Dict, Any, Optional
+import operator
 
-# LangChain imports
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
-from langchain_openai import ChatOpenAI
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.tools.wikipedia.tool import WikipediaQueryRun
 from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
 from langchain_community.tools.arxiv.tool import ArxivQueryRun
+from langchain_community.utilities.arxiv import ArxivAPIWrapper
+
 from langgraph.graph import StateGraph, END
-from langgraph.prebuilt import ToolNode, tools_condition
-from typing import TypedDict
+from langgraph.prebuilt import ToolInvocation
+from langchain_core.messages import BaseMessage, FunctionMessage, HumanMessage, AIMessage
+from langchain_openai import ChatOpenAI
 
+# --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-class MessagesState(TypedDict):
-    messages: List[BaseMessage]
-
-try:
-    with open("system_prompt.txt", "r", encoding="utf-8") as f:
-        system_prompt = f.read()
-        print("Successfully loaded system prompt from file")
-except FileNotFoundError:
-    system_prompt = (
-        "You are a helpful AI assistant that uses tools to find information and answer questions.\n"
-        "When you don't know something, use the available tools to look up information. Be concise, direct, and provide accurate responses.\n"
-        "Always cite your sources when using information from searches or reference materials."
-    )
-    print("System prompt file not found, using default prompt")
+# --- Environment Setup ---
+load_dotenv()
+
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") # Assuming Tavily might also need an API key
+
+if not OPENROUTER_API_KEY:
+    print("Warning: OPENROUTER_API_KEY not found in .env file. The LLM will not function.")
 
-class AdvancedAgent:
+# --- Tool Setup ---
+tavily_tool = TavilySearchResults(max_results=3, api_key=TAVILY_API_KEY if TAVILY_API_KEY else "placeholder_tavily_key") # Add placeholder if not found
+wikipedia_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper(top_k_results=2, doc_content_chars_max=2000))
+arxiv_tool = ArxivQueryRun(api_wrapper=ArxivAPIWrapper(top_k_results=2, doc_content_chars_max=2000))
+
+tools = [tavily_tool, wikipedia_tool, arxiv_tool]
+
+# --- LangGraph Agent Definition ---
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+    next_action: Optional[str] # To decide if we need to call tools or respond
+
+class LangGraphAgent:
     def __init__(self):
-        print("Initializing AdvancedAgent with LangGraph, Wikipedia, Arxiv, and LLM via OpenRouter")
-        load_dotenv()
-        
-        # Check if OpenRouter API key is set
-        if not os.getenv("OPENROUTER_API_KEY"):
-            print("WARNING: OPENROUTER_API_KEY environment variable not set!")
-            print("Please set this in your HF Space secrets or .env file")
-            
-        try:
-            self.graph = self.build_graph()
-            print("Graph successfully built")
-        except Exception as e:
-            print(f"Error building agent graph: {e}")
-            traceback.print_exc()
-            raise
-        
-    def build_graph(self):
-        # Add consistent headers for OpenRouter
-        headers = {
-            "HTTP-Referer": "https://huggingface.co/",
-            "X-Title": "HF Agent"
-        }
-        
-    try:
-        # Try alternative OpenAI model if Gemini is causing issues
-        model_name = "google/gemini-2.0-flash-001"
-        # Fallback to OpenAI if needed
-        # model_name = "gpt-3.5-turbo" 
-        
-        llm = ChatOpenAI(
-            model=model_name,
-            temperature=0,
-            openai_api_key=os.getenv("OPENROUTER_API_KEY"),
-            openai_api_base="https://openrouter.ai/api/v1",
-            headers=headers
-        )
-        print(f"LLM initialized: {model_name} via OpenRouter")
-        
-        wikipedia_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
-        arxiv_tool = ArxivQueryRun()
-        tavily_search = TavilySearchResults(max_results=5)
-        tools = [wikipedia_tool, arxiv_tool, tavily_search]
-        print(f"Initialized {len(tools)} tools: Wikipedia, Arxiv, Tavily Search")
-        
-        llm_with_tools = llm.bind_tools(tools)
-
-        def assistant(state: MessagesState):
-            messages = state["messages"]
-            response = llm_with_tools.invoke(messages)
-            return {"messages": messages + [response]}  # Always return dict
-
-        tools_node = ToolNode(tools)
-
-        builder = StateGraph(MessagesState)
-        builder.add_node("assistant", assistant)
-        builder.add_node("tools", tools_node)
-        builder.add_edge("tools", "assistant")
-        builder.add_edge("assistant", "tools")
-        builder.set_entry_point("assistant")
-        builder.add_conditional_edges(
-            "tools",
-            lambda x: "assistant",
-            {"assistant": "assistant"}
+        print("LangGraphAgent initializing...")
+        if not OPENROUTER_API_KEY:
+            raise ValueError("OPENROUTER_API_KEY is not set. Cannot initialize LLM.")
+
+        self.llm = ChatOpenAI(
+            model="google/gemini-2.0-flash-001",
+            api_key=OPENROUTER_API_KEY,
+            base_url="https://openrouter.ai/api/v1"
         )
-        builder.add_conditional_edges(
-            "assistant",
-            tools_condition,
-            {"tools": "tools", END: END}
+        self.tools_map = {tool.name: tool for tool in tools}
+        self.graph = self._build_graph()
+        print("LangGraphAgent initialized.")
+
+    def _build_graph(self):
+        workflow = StateGraph(AgentState)
+
+        workflow.add_node("llm", self._call_llm)
+        workflow.add_node("tools", self._tool_node)
+
+        workflow.set_entry_point("llm")
+
+        workflow.add_conditional_edges(
+            "llm",
+            self._should_call_tools,
+            {
+                "continue": "tools",
+                "end": END
+            }
         )
-        return builder.compile()
-    
+        workflow.add_edge("tools", "llm")
+        return workflow.compile()
+
+    def _should_call_tools(self, state: AgentState) -> str:
+        print("LLM deciding next step...")
+        last_message = state["messages"][-1]
+        if hasattr(last_message, "tool_calls") and last_message.tool_calls:
+            print(f"LLM decided to call tools: {last_message.tool_calls}")
+            return "continue"
+        print("LLM decided to end.")
+        return "end"
+
+    def _call_llm(self, state: AgentState) -> Dict[str, Any]:
+        print("Calling LLM...")
+        # Bind tools to LLM for function calling
+        llm_with_tools = self.llm.bind_tools(tools)
+        response = llm_with_tools.invoke(state["messages"])
+        print(f"LLM response: {response.content[:100]}...")
+        return {"messages": [response]}
+
+    def _tool_node(self, state: AgentState) -> Dict[str, Any]:
+        print("Executing tools...")
+        tool_messages = []
+        last_message = state["messages"][-1]
+
+        if not hasattr(last_message, "tool_calls") or not last_message.tool_calls:
+            print("No tool calls found in the last message.")
+            # This case should ideally be handled by the conditional edge, but as a fallback:
+            return {"messages": [AIMessage(content="No tools to call, proceeding.")]}
+
+        for tool_call in last_message.tool_calls:
+            tool_name = tool_call["name"]
+            tool_args = tool_call["args"]
+            print(f"Calling tool: {tool_name} with args: {tool_args}")
+            if tool_name in self.tools_map:
+                try:
+                    tool_result = self.tools_map[tool_name].invoke(tool_args)
+                    print(f"Tool {tool_name} result (first 100 chars): {str(tool_result)[:100]}...")
+                    tool_messages.append(FunctionMessage(content=str(tool_result), name=tool_name, tool_call_id=tool_call["id"]))
+                except Exception as e:
+                    print(f"Error executing tool {tool_name}: {e}")
+                    tool_messages.append(FunctionMessage(content=f"Error executing tool {tool_name}: {e}", name=tool_name, tool_call_id=tool_call["id"]))
+            else:
+                print(f"Tool {tool_name} not found.")
+                tool_messages.append(FunctionMessage(content=f"Tool {tool_name} not found.", name=tool_name, tool_call_id=tool_call["id"]))
+        return {"messages": tool_messages}
+
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        messages = [
-            SystemMessage(content=system_prompt),
-            HumanMessage(content=question)
-        ]
-        
-        # Print message details for debugging
-        print(f"Sending {len(messages)} messages to the LLM")
+        initial_state = {"messages": [HumanMessage(content=question)]}
         
+        # The GAIA prompt example suggests not including "FINAL ANSWER" and just replying with the answer.
+        # We need to ensure the LLM is prompted to provide a direct answer after tool use.
+        # For simplicity in this template, we will take the last AI message content as the answer.
+        # A more robust solution might involve a specific "final answer" node or prompt engineering.
+
+        final_graph_state = None
         try:
-            # Initial state must be a dict with "messages" key!
-            print("Invoking agent graph...")
-            result = self.graph.invoke({"messages": messages})
-            print("Agent graph execution completed")
-            
-            final_messages = result["messages"]
-            print(f"Got {len(final_messages)} messages in result")
-            
-            # Extract only AI messages from the final state
-            ai_messages = [msg for msg in final_messages if isinstance(msg, AIMessage)]
-            print(f"Found {len(ai_messages)} AI messages in the final state")
-            
-            if not ai_messages:
-                print("No AI messages found in the final state")
-                return "I wasn't able to generate a proper response. Please try again."
-            
-            # Get the last AI message - this should be the final answer
-            last_message = ai_messages[-1]
-            
-            # Return the content of the last AI message
-            if last_message.content:
-                return last_message.content
+            for event in self.graph.stream(initial_state, {"recursion_limit": 10}): # Added recursion limit
+                # print(f"Graph event: {event}") # For debugging stream
+                if END in event:
+                    final_graph_state = event[END]
+                    break
+                # Update final_graph_state with the latest state from any node
+                # This ensures we have the latest messages even if END is not directly reached by llm
+                # (e.g. if recursion limit is hit)
+                for key in event:
+                    if key != END:
+                        final_graph_state = event[key]
+
+            if final_graph_state and final_graph_state["messages"]:
+                # Get the last AI message as the answer
+                for msg in reversed(final_graph_state["messages"]):
+                    if isinstance(msg, AIMessage) and not msg.tool_calls:
+                        answer = msg.content.strip()
+                        # Ensure no "FINAL ANSWER:" prefix as per GAIA instructions
+                        if answer.upper().startswith("FINAL ANSWER:"):
+                            answer = answer[len("FINAL ANSWER:"):].strip()
+                        print(f"Agent returning answer: {answer}")
+                        return answer
+                # Fallback if no suitable AI message is found
+                print("No suitable AI message found for final answer. Returning last message content.")
+                # This might be a tool call or an intermediate thought, not ideal.
+                return str(final_graph_state["messages"][-1].content) if final_graph_state["messages"] else "Error: No messages in final state."
             else:
-                print("Last AI message has empty content")
-                return "I wasn't able to generate a proper response. Please try again."
-            
+                print("Error: Agent did not reach a final state or no messages found.")
+                return "Error: Agent did not produce a conclusive answer."
+
         except Exception as e:
-            print(f"Error running agent graph: {e}")
+            print(f"Error during agent execution: {e}")
+            import traceback
             traceback.print_exc()
-            return f"Sorry, I encountered an error while processing your question: {str(e)}"
+            return f"Error during agent execution: {e}"
 
+# --- Main Evaluation Logic (Modified from starter) ---
 def run_and_submit_all(profile: gr.OAuthProfile | None):
-    space_id = os.getenv("SPACE_ID", "i-dhilip/Final_Assignment_Template")  # Default to your space ID
+    """
+    Fetches all questions, runs the LangGraphAgent on them, submits all answers,
+    and displays the results.
+    """
+    space_id = os.getenv("SPACE_ID")
+
     if profile:
         username = f"{profile.username}"
         print(f"User logged in: {username}")
@@ -159,41 +182,28 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         print("User not logged in.")
         return "Please Login to Hugging Face with the button.", None
 
+    if not OPENROUTER_API_KEY:
+         return "Error: OPENROUTER_API_KEY not found. Please set it in your .env file.", None
+    if not TAVILY_API_KEY:
+        print("Warning: TAVILY_API_KEY not found. Tavily search might not work as expected.")
+        # return "Error: TAVILY_API_KEY not found. Please set it in your .env file.", None
+
     api_url = DEFAULT_API_URL
     questions_url = f"{api_url}/questions"
     submit_url = f"{api_url}/submit"
-    
-    print(f"Checking OpenRouter API key...")
-    if not os.getenv("OPENROUTER_API_KEY"):
-        return "Error: OPENROUTER_API_KEY not set. Please add this to your Space secrets.", None
 
     try:
-        agent = AdvancedAgent()
-    except Exception as e:
-        error_msg = f"Error instantiating agent: {e}"
-        print(error_msg)
-        traceback.print_exc()
-        return error_msg, None
-
-    # Generate a simple test question to verify the agent works
-    test_question = "What is the capital of France?"
-    print(f"Testing agent with a simple question: '{test_question}'")
-    try:
-        test_response = agent(test_question)
-        print(f"Test response: {test_response[:100]}...")
-        if "I wasn't able to generate" in test_response or "error" in test_response.lower():
-            print("WARNING: Agent test response indicates potential issues")
+        agent = LangGraphAgent()
     except Exception as e:
-        print(f"Agent test failed: {e}")
-        traceback.print_exc()
-        return f"Agent test failed: {e}", None
+        print(f"Error instantiating agent: {e}")
+        return f"Error initializing agent: {e}", None
 
-    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
-    print(agent_code)
+    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else "local_run_no_space_id"
+    print(f"Agent code link: {agent_code}")
 
     print(f"Fetching questions from: {questions_url}")
     try:
-        response = requests.get(questions_url, timeout=15)
+        response = requests.get(questions_url, timeout=20)
         response.raise_for_status()
         questions_data = response.json()
         if not questions_data:
@@ -207,9 +217,6 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         print(f"Error decoding JSON response from questions endpoint: {e}")
         print(f"Response text: {response.text[:500]}")
         return f"Error decoding server response for questions: {e}", None
-    except Exception as e:
-        print(f"An unexpected error occurred fetching questions: {e}")
-        return f"An unexpected error occurred fetching questions: {e}", None
 
     results_log = []
     answers_payload = []
@@ -221,6 +228,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
+            print(f"\n--- Processing Task ID: {task_id} ---")
             submitted_answer = agent(question_text)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
@@ -233,7 +241,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
 
     submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
-    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
+    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user 	'{username}	'..."
     print(status_update)
 
     print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
@@ -278,22 +286,32 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         results_df = pd.DataFrame(results_log)
         return status_message, results_df
 
+# --- Gradio Interface (Mostly as provided) ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Advanced Agent Evaluation Runner")
+    gr.Markdown("# LangGraph GAIA Agent Evaluation Runner")
     gr.Markdown(
         """
         **Instructions:**
-        1.  This space implements an advanced agent using LangGraph with Wikipedia, Arxiv, and Tavily Search tools, powered by Gemini 2.0 Flash LLM.
-        2.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
-        3.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
+        1.  **Clone this space** if you haven't already.
+        2.  **Create a `.env` file** in the root of your space with your API keys:
+            ```
+            OPENROUTER_API_KEY="your_openrouter_api_key"
+            TAVILY_API_KEY="your_tavily_api_key" # Optional, but recommended for TavilySearch tool
+            ```
+        3.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
+        4.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
         ---
-        **Note:** 
-        The evaluation might take some time as the agent processes all questions through the tools.
+        **Disclaimers:**
+        -   Ensure your Hugging Face Space is public for the `agent_code` link to be verifiable.
+        -   Submitting all answers can take some time as the agent processes each question.
+        -   This agent uses LangGraph with TavilySearch, Wikipedia, Arxiv, and Google Gemini via OpenRouter.
         """
     )
 
     gr.LoginButton()
+
     run_button = gr.Button("Run Evaluation & Submit All Answers")
+
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
     results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
 
@@ -306,17 +324,22 @@ if __name__ == "__main__":
     print("\n" + "-"*30 + " App Starting " + "-"*30)
     space_host_startup = os.getenv("SPACE_HOST")
     space_id_startup = os.getenv("SPACE_ID")
+
     if space_host_startup:
         print(f"✅ SPACE_HOST found: {space_host_startup}")
         print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
     else:
         print("ℹ️  SPACE_HOST environment variable not found (running locally?).")
+
     if space_id_startup:
         print(f"✅ SPACE_ID found: {space_id_startup}")
         print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
         print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
     else:
         print("ℹ️  SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")
+
     print("-"*(60 + len(" App Starting ")) + "\n")
-    print("Launching Gradio Interface for Advanced Agent Evaluation...")
-    demo.launch(debug=True, share=False)
+
+    print("Launching Gradio Interface for LangGraph GAIA Agent Evaluation...")
+    demo.launch(debug=True, share=False)
+