Update app.py

app.py

@@ -1,674 +1,267 @@
-"""LangGraph Agent with Gradio Interface"""
 import os
+import gradio as gr
+import requests
+import inspect
+import pandas as pd
 from dotenv import load_dotenv
-from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import tools_condition
-from langgraph.prebuilt import ToolNode
+from typing import List, Dict, Any, Tuple, Optional
+
+# LangChain imports
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+from langchain_core.messages import BaseMessage
+from langchain.schema import Document
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_groq import ChatGroq
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.document_loaders import ArxivLoader
-from langchain_community.vectorstores import SupabaseVectorStore
-from langchain_core.messages import SystemMessage, HumanMessage
-from langchain_core.tools import tool
-from langchain.tools.retriever import create_retriever_tool
-from supabase.client import Client, create_client
-import json
-import requests
-import google.generativeai as genai
-import gradio as gr
-from typing import Literal
-from typing import TypedDict, Annotated, Sequence
-import operator
-
-# Load environment variables
-load_dotenv()
-
-# Setup API keys
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
-HF_API_TOKEN = os.getenv("HF_API_TOKEN")
-OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
-
-if GOOGLE_API_KEY:
-    genai.configure(api_key=GOOGLE_API_KEY)
-else:
-    print("Warning: GOOGLE_API_KEY not found")
-
-
-@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 wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
-    
-    Args:
-        query: The search query."""
-    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    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 {"wiki_results": formatted_search_docs}
-
-@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}
-
-@tool
-def arvix_search(query: str) -> str:
-    """Search Arxiv for a query and return maximum 3 result.
-    
-    Args:
-        query: The search query."""
-    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"arvix_results": formatted_search_docs}
-
-
-
-# load the system prompt from the file
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    system_prompt = f.read()
-
-# System message
-sys_msg = SystemMessage(content=system_prompt)
-
-# build a retriever
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") #  dim=768
-supabase: Client = create_client(
-    os.environ.get("SUPABASE_URL"), 
-    os.environ.get("SUPABASE_SERVICE_KEY"))
-vector_store = SupabaseVectorStore(
-    client=supabase,
-    embedding= embeddings,
-    table_name="documents",
-    query_name="match_documents_langchain",
-)
-create_retriever_tool = create_retriever_tool(
-    retriever=vector_store.as_retriever(),
-    name="Question Search",
-    description="A tool to retrieve similar questions from a vector store.",
-)
-
-tools = [
-    multiply,
-    add,
-    subtract,
-    divide,
-    modulus,
-    wiki_search,
-    web_search,
-    arvix_search,
-]
-
-# Build graph function
-def build_graph(provider: str = "google"):
-    """Build the graph"""
-    # Load environment variables from .env file
-    if provider == "google":
-        # Google Gemini
+from langchain_community.tools.wikipedia.tool import WikipediaQueryRun
+from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
+from langchain_community.tools.arxiv.tool import ArxivQueryRun
+from langgraph.graph import StateGraph, END
+from langgraph.graph.nodes.tools import ToolNode
+from langgraph.prebuilt import ToolInvocation, tools_condition
+from langgraph.prebuilt.tool_executor import ToolExecutor
+from dataclasses import dataclass
+from typing import TypedDict, List, Annotated, Literal
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+
+# Constants
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+
+# Define the state for the agent
+class MessagesState(TypedDict):
+    messages: List[BaseMessage]
+
+# Load system prompt
+try:
+    with open("system_prompt.txt", "r", encoding="utf-8") as f:
+        system_prompt = f.read()
+except FileNotFoundError:
+    system_prompt = """You are a helpful AI assistant that uses tools to find information and answer questions.
+When you don't know something, use the available tools to look up information. Be concise, direct, and provide accurate responses.
+Always cite your sources when using information from searches or reference materials."""
+
+# Advanced agent using LangGraph
+class AdvancedAgent:
+    def __init__(self):
+        print("Initializing AdvancedAgent with LangGraph, Wikipedia, Arxiv, and Gemini 2.0 Flash")
+        load_dotenv()  # Load environment variables from .env file
+        
+        # Initialize the graph
+        self.graph = self.build_graph()
+        print("Graph successfully built")
+        
+    def build_graph(self):
+        """Build the LangGraph agent with necessary tools"""
+        # Initialize LLM
         llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    elif provider == "huggingface":
-        # TODO: Add huggingface endpoint
-        llm = ChatHuggingFace(
-            llm=HuggingFaceEndpoint(
-                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
-                temperature=0,
-            ),
-        )
-    else:
-        raise ValueError("Invalid provider. Choose 'google', or 'huggingface'.")
-    # Bind tools to LLM
-    llm_with_tools = llm.bind_tools(tools)
-
-    # Node
-    def assistant(state: MessagesState):
-        """Assistant node"""
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
-    def retriever(state: MessagesState):
-        """Retriever node"""
-        similar_question = vector_store.similarity_search(state["messages"][0].content)
-        example_msg = HumanMessage(
-            content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
-        )
-        return {"messages": [sys_msg] + state["messages"] + [example_msg]}
-
-    builder = StateGraph(MessagesState)
-    builder.add_node("retriever", retriever)
-    builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
-    builder.add_edge(START, "retriever")
-    builder.add_edge("retriever", "assistant")
-    builder.add_conditional_edges(
-        "assistant",
-        tools_condition,
-    )
-    builder.add_edge("tools", "assistant")
-
-    # Compile graph
-    return builder.compile()
-
-
-
-# Tool Definitions
-# @tool
-# def multiply(a: int, b: int) -> int:
-#     """Multiply two numbers."""
-#     return a * b
-
-# @tool
-# def add(a: int, b: int) -> int:
-#     """Add two numbers."""
-#     return a + b
-
-# @tool
-# def modulus(a: int, b: int) -> int:
-#     """Get the modulus of two numbers."""
-#     return a % b
-
-# @tool
-# def subtract(a: int, b: int) -> int:
-#     """Subtract two numbers."""
-#     return a - b
-
-# @tool
-# def divide(a: int, b: int) -> int:
-#     """Divide two numbers."""
-#     if b == 0:
-#         raise ValueError("Cannot divide by zero.")
-#     return a / b
-
-# @tool
-# def arvix_search(query: str) -> str:
-#     """Search Arxiv for a query and return maximum 3 results."""
-#     try:
-#         search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-#         formatted_search_docs = "\n\n---\n\n".join(
-#             [f'<Document source="{doc.metadata["source"]}"/>\n{doc.page_content[:1000]}\n</Document>' 
-#              for doc in search_docs])
-#         return {"arvix_results": formatted_search_docs}
-#     except Exception as e:
-#         return {"arvix_results": f"Error: {str(e)}"}
-
-# @tool
-# def execute_python(code: str) -> str:
-#     """Execute Python code securely and return results. Handles calculations and data analysis."""
-#     try:
-#         parsed = ast.parse(code)
-#         if any(isinstance(node, (ast.Import, ast.ImportFrom)) for node in parsed.body):
-#             return "Cannot import modules for security reasons"
+        print("LLM initialized: Gemini 2.0 Flash")
         
-#         restricted_builtins = {'open', 'eval', 'exec', '__import__'}
-#         for node in ast.walk(parsed):
-#             if isinstance(node, ast.Name) and node.id in restricted_builtins:
-#                 return "Restricted function used"
-                
-#         return str(eval(code, {"__builtins__": {}}, {}))
-#     except Exception as e:
-#         return f"Execution error: {str(e)}"
-
-# @tool
-# def process_file(file_path: str) -> Dict[str, Any]:
-#     """Process uploaded files (Excel, CSV, TXT) and extract structured data."""
-#     try:
-#         if file_path.endswith(('.xlsx', '.xls')):
-#             df = pd.read_excel(file_path)
-#             return {"data": df.head(10).to_dict(), "summary": df.describe().to_dict()}
-#         elif file_path.endswith('.csv'):
-#             df = pd.read_csv(file_path)
-#             return {"data": df.head(10).to_dict(), "summary": df.describe().to_dict()}
-#         elif file_path.endswith('.txt'):
-#             with open(file_path, 'r') as f:
-#                 content = f.read(2000)
-#             return {"content": content}
-#         return {"error": "Unsupported file format"}
-#     except Exception as e:
-#         return {"error": str(e)}
-
-# @tool
-# def enhanced_wiki_search(query: str) -> str:
-#     """Access full Wikipedia articles with history and specific versions."""
-#     try:
-#         loader = WikipediaLoader(query=query, load_max_docs=2, doc_content_chars_max=4000)
-#         docs = loader.load()
-#         return "\n\n".join([
-#             f"Title: {doc.metadata['title']}\n"
-#             f"URL: {doc.metadata['source']}\n"
-#             f"Content: {doc.page_content[:3000]}..."
-#             for doc in docs
-#         ])
-#     except Exception as e:
-#         return f"Wikipedia Error: {str(e)}"
-
-# @tool
-# def media_search(query: str) -> str:
-#     """Specialized search for media content (videos, images, celebrities)."""
-#     try:
-#         search = DuckDuckGoSearchResults(max_results=5)
-#         results = search.invoke(f"site:imdb.com OR site:youtube.com {query}")
-#         return "\n".join([
-#             f"Source: {res['link']}\nSnippet: {res['snippet']}"
-#             for res in results[:3]
-#         ])
-#     except Exception as e:
-#         return f"Media Search Error: {str(e)}"
-
-# @tool
-# def academic_search(query: str) -> str:
-#     """Search academic databases and educational resources."""
-#     try:
-#         arxiv_docs = ArxivLoader(query=query, load_max_docs=2).load()
-#         web_docs = DuckDuckGoSearchResults(max_results=3).invoke(f"filetype:pdf {query}")
-#         return f"Arxiv Results:\n{arxiv_docs[0].page_content[:1000]}\n\nWeb Results:\n{web_docs[0]['snippet']}"
-#     except Exception as e:
-#         return f"Academic Search Error: {str(e)}"
-
-
-# @tool
-# def web_search(query: str) -> str:
-#     """Search DuckDuckGo for a query and return maximum 3 results."""
-#     try:
-#         search = DuckDuckGoSearchResults(max_results=3)
-#         search_docs = search.invoke(query)
-#         formatted_search_docs = "\n\n---\n\n".join(
-#             [f'<Document source="{doc["link"]}"/>\n{doc["snippet"]}\n</Document>' 
-#              for doc in search_docs])
-#         return {"web_results": formatted_search_docs}
-#     except Exception as e:
-#         return {"web_results": f"Error: {str(e)}"}
-
-# @tool
-# def summarize_text(text: str) -> str:
-#     """Summarize a long text into key points."""
-#     if not text or len(text) < 100:
-#         return "Text too short to summarize."
-#     return f"Summary of the provided text with key points."
-
-# @tool
-# def parse_query(query: str) -> dict:
-#     """Parse a complex query into its key components for better search."""
-#     parts = query.split()
-#     return {
-#         "main_topic": parts[0] if parts else "",
-#         "subtopics": parts[1:3] if len(parts) > 1 else [],
-#         "context": " ".join(parts[3:]) if len(parts) > 3 else ""
-#     }
-
-# # System Prompt Setup
-# system_prompt = """You are a POWERFUL assistant REQUIRED to answer ALL questions using available tools.
-# STRICT RULES:
-# 1. NEVER say you can't answer - ALWAYS use tools
-# 2. Combine information from multiple tools when needed
-# 3. For calculations, use execute_python
-# 4. For files, use process_file
-# 5. For media/celebrities, use media_search
-# 6. For academic content, use academic_search
-# 7. ALWAYS format final answer as: FINAL ANSWER: [your answer]
-
-# AVAILABLE TOOLS:
-# - execute_python: Math/code execution
-# - process_file: Analyze uploaded files
-# - enhanced_wiki_search: Full Wikipedia access
-# - media_search: Videos/images/celebrities
-# - academic_search: Textbooks/papers
-# - web_search: General web search
-# - vector_store: Previous knowledge
-
-# YOU MUST USE THESE TOOLS TO ANSWER ALL QUESTIONS!"""
-# sys_msg = SystemMessage(content=system_prompt)
-
-# # Vector Store Setup
-# try:
-#     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-#     vector_store = Chroma(
-#         collection_name="documents",
-#         embedding_function=embeddings,
-#         persist_directory="./chroma_db"
-#     )
-# except Exception as e:
-#     print(f"Error initializing vector store: {e}")
-#     vector_store = None
-
-# # Tool Configuration
-# tools = [
-#     multiply, add, subtract, divide, modulus,
-#     enhanced_wiki_search, media_search, web_search, arvix_search,
-#     academic_search, summarize_text, parse_query, DuckDuckGoSearchResults(max_results=5)
-# ]
-
-# if vector_store:
-#     tools.append(
-#         create_retriever_tool(
-#             vector_store.as_retriever(),
-#             name="Question Search",
-#             description="Retrieves similar questions from vector store"
-#         )
-#     )
-
-# # Model Configuration
-# MODEL_REGISTRY = {
-#     "gemini-2.0-flash": {
-#         "provider": "google",
-#         "model": "gemini-2.0-flash",
-#         "temperature": 0.2,
-#         "max_tokens": 2048
-#     },
-#     "gemini-1.5-flash": {
-#         "provider": "google",
-#         "model": "gemini-1.5-flash", 
-#         "temperature": 0.2,
-#         "max_tokens": 2048
-#     },
-#     "kimi-vl-a3b-thinking": {
-#         "provider": "openrouter",
-#         "model": "moonshotai/kimi-vl-a3b-thinking:free",
-#         "temperature": 0.2,
-#         "max_tokens": 2048
-#     }
-# }
-
-# def get_llm(model_name: str = "gemini-2.0-flash"):
-#     """Initialize LLM with error handling"""
-#     config = MODEL_REGISTRY.get(model_name, MODEL_REGISTRY["gemini-2.0-flash"])
-#     provider = config.get("provider", "google")
-    
-#     try:
-#         if provider == "google":
-#             if not GOOGLE_API_KEY:
-#                 print(f"Error initializing {model_name}: GOOGLE_API_KEY not found")
-#                 return None
-#             return ChatGoogleGenerativeAI(
-#                 model=config["model"],
-#                 temperature=config["temperature"],
-#                 max_output_tokens=config["max_tokens"],
-#                 convert_system_message_to_human=True
-#             )
-#         elif provider == "openrouter":
-#             if not OPENROUTER_API_KEY:
-#                 print(f"Error initializing {model_name}: OPENROUTER_API_KEY not found")
-#                 return None
-#             return ChatOpenAI(
-#                 model=config["model"],
-#                 temperature=config["temperature"],
-#                 max_tokens=config["max_tokens"],
-#                 openai_api_key=OPENROUTER_API_KEY,
-#                 openai_api_base="https://openrouter.ai/api/v1",
-#                 model_kwargs={
-#                     "headers": {
-#                         "HTTP-Referer": "https://your-site.com",
-#                         "X-Title": "Agent Evaluation"
-#                     }
-#                 }
-#             )
-#         else:
-#             print(f"Unknown provider {provider} for model {model_name}")
-#             return None
-#     except Exception as e:
-#         print(f"Error initializing {model_name}: {e}")
-#         return None
-
-# # Graph Builder
-# def build_graph():
-#     """Build LangGraph agent workflow"""
-#     primary_llm = get_llm("gemini-2.0-flash")
-#     fallback_llm = get_llm("gemini-1.5-flash")
-#     kimi_llm = get_llm("kimi-vl-a3b-thinking")
-    
-#     llms = [llm for llm in [primary_llm, fallback_llm, kimi_llm] if llm is not None]
-    
-#     if not llms:
-#         raise RuntimeError("Failed to initialize any LLM")
-
-#     current_llm_index = 0
-
-#     def assistant(state: MessagesState):
-#         nonlocal current_llm_index
-#         for attempt in range(len(llms)):
-#             try:
-#                 llm = llms[current_llm_index]
-#                 llm_with_tools = llm.bind_tools(tools)
-                
-#                 messages = state["messages"].copy()
-#                 if len(messages) > 0 and isinstance(messages[0], HumanMessage):
-#                     tool_instruction = HumanMessage(content="Use available tools to answer.")
-#                     messages.append(tool_instruction)
-                
-#                 response = llm_with_tools.invoke(messages)
-#                 current_llm_index = (current_llm_index + 1) % len(llms)
-#                 return {"messages": [response]}
-#             except Exception as e:
-#                 print(f"Model {llms[current_llm_index]} failed: {e}")
-#                 current_llm_index = (current_llm_index + 1) % len(llms)
-#                 if attempt == len(llms) - 1:
-#                     error_msg = HumanMessage(content=f"All models failed: {str(e)}")
-#                     return {"messages": [error_msg]}
-
-#     def retriever(state: MessagesState):
-#         try:
-#             if vector_store:
-#                 similar_questions = vector_store.similarity_search(
-#                     state["messages"][0].content,
-#                     k=1
-#                 )
-#                 example_content = "Similar question reference: \n\n" + \
-#                     (similar_questions[0].page_content if similar_questions 
-#                      else "No similar questions found")
-#             else:
-#                 example_content = "Vector store not available"
-            
-#             return {"messages": [sys_msg] + state["messages"] + [HumanMessage(content=example_content)]}
-#         except Exception as e:
-#             error_msg = HumanMessage(content=f"Retrieval error: {str(e)}")
-#             return {"messages": [error_msg]}
-
-#     builder = StateGraph(MessagesState)
-#     builder.add_node("retriever", retriever)
-#     builder.add_node("assistant", assistant)
-#     builder.add_node("tools", ToolNode(tools))
-    
-#     builder.add_edge(START, "retriever")
-#     builder.add_edge("retriever", "assistant")
-#     builder.add_conditional_edges("assistant", tools_condition)
-#     builder.add_edge("tools", "assistant")
+        # Initialize tools
+        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")
+        
+        # Create tool executor
+        tool_executor = ToolExecutor(tools)
+        
+        # System message
+        sys_msg = SystemMessage(content=system_prompt)
+        
+        # Bind tools to LLM
+        llm_with_tools = llm.bind_tools(tools)
+        
+        # Define nodes
+        def assistant(state: MessagesState):
+            """Assistant node that processes messages and generates responses"""
+            messages = state["messages"]
+            response = llm_with_tools.invoke(messages)
+            return {"messages": state["messages"] + [response]}
+        
+        def tools_node(state: MessagesState, tool_calls: List[ToolInvocation]):
+            """Execute tool calls and return results"""
+            results = []
+            for tool_call in tool_calls:
+                result = tool_executor.invoke(tool_call)
+                msg = AIMessage(content=str(result), tool_call_id=tool_call.id)
+                results.append(msg)
+            return {"messages": state["messages"] + results}
+        
+        # Build the graph
+        builder = StateGraph(MessagesState)
+        
+        # Add nodes
+        builder.add_node("assistant", assistant)
+        builder.add_node("tools", tools_node)
+        
+        # Add edges
+        builder.add_edge("assistant", "tools", condition=tools_condition)
+        builder.add_edge("tools", "assistant")
+        builder.add_edge("assistant", END, condition=lambda state: not tools_condition(state))
+        
+        # Compile graph
+        return builder.compile()
     
-#     return builder.compile()
-
-# Agent Class
-class BasicAgent:
-    def __init__(self):
-        self.graph = build_graph()
-
     def __call__(self, question: str) -> str:
+        """Process a question through the agent graph and return the response"""
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        
+        # Create initial state with system message and human question
+        messages = [
+            SystemMessage(content=system_prompt),
+            HumanMessage(content=question)
+        ]
+        
+        # Run the graph
         try:
-            messages = [HumanMessage(content=question)]
             result = self.graph.invoke({"messages": messages})
-            last_message = result['messages'][-1].content
+            # Extract the last AI message as the answer
+            for msg in reversed(result["messages"]):
+                if isinstance(msg, AIMessage) and not getattr(msg, "tool_call_id", None):
+                    return msg.content
             
-            if "FINAL ANSWER: " in last_message:
-                answer_part = last_message.split("FINAL ANSWER: ")[-1].strip()
-                return answer_part[:-2].strip() if answer_part.endswith('"}') else answer_part
-            elif "Answer:" in last_message:
-                answer_part = last_message.split("Answer:")[-1].strip()
-                return answer_part[:-2].strip() if answer_part.endswith('"}') else answer_part
-            return last_message
+            # Fallback if no valid AI message found
+            return "I wasn't able to generate a proper response. Please try again."
         except Exception as e:
-            return f"Agent error: {str(e)}"
+            print(f"Error running agent graph: {e}")
+            return f"Sorry, I encountered an error while processing your question: {str(e)}"
 
-
-
-
-# Updated Agent Class
-# class BasicAgent:
-#     """LangGraph Agent Interface"""
-#     def __init__(self):
-#         self.graph = build_graph()
-
-#     def __call__(self, question: str) -> str:
-#         try:
-#             messages = [HumanMessage(content=question)]
-#             result = self.graph.invoke({"messages": messages})
-#             last_message = result['messages'][-1].content
-            
-#             # Improved content extraction
-#             if "FINAL ANSWER: " in last_message:
-#                 return last_message.split("FINAL ANSWER: ")[-1].strip()
-#             elif "Answer:" in last_message:
-#                 return last_message.split("Answer:")[-1].strip()
-#             return last_message
-#         except Exception as e:
-#             return f"Agent processing error: {str(e)}"
-
-
-
-
-
-# Gradio Interface Functions
+# Function to run and submit all questions
 def run_and_submit_all(profile: gr.OAuthProfile | None):
-    """Evaluation runner function"""
-    if not profile:
+    """
+    Fetches all questions, runs the AdvancedAgent on them, submits all answers,
+    and displays the results.
+    """
+    # --- Determine HF Space Runtime URL and Repo URL ---
+    space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code
+
+    if profile:
+        username= f"{profile.username}"
+        print(f"User logged in: {username}")
+    else:
+        print("User not logged in.")
         return "Please Login to Hugging Face with the button.", None
 
-    space_id = os.getenv("SPACE_ID")
-    api_url = "https://agents-course-unit4-scoring.hf.space"
-    username = profile.username
-    results_log = []
+    api_url = DEFAULT_API_URL
+    questions_url = f"{api_url}/questions"
+    submit_url = f"{api_url}/submit"
 
+    # 1. Instantiate Agent
     try:
-        agent = BasicAgent()
-        agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
-        
-        # Fetch questions
-        response = requests.get(f"{api_url}/questions", timeout=15)
+        agent = AdvancedAgent()
+    except Exception as e:
+        print(f"Error instantiating agent: {e}")
+        return f"Error initializing agent: {e}", None
+    
+    # In the case of an app running as a hugging Face space, this link points toward your codebase
+    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
+    print(agent_code)
+
+    # 2. Fetch Questions
+    print(f"Fetching questions from: {questions_url}")
+    try:
+        response = requests.get(questions_url, timeout=15)
         response.raise_for_status()
         questions_data = response.json()
+        if not questions_data:
+             print("Fetched questions list is empty.")
+             return "Fetched questions list is empty or invalid format.", None
+        print(f"Fetched {len(questions_data)} questions.")
+    except requests.exceptions.RequestException as e:
+        print(f"Error fetching questions: {e}")
+        return f"Error fetching questions: {e}", None
+    except requests.exceptions.JSONDecodeError as e:
+         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
+
+    # 3. Run your Agent
+    results_log = []
+    answers_payload = []
+    print(f"Running agent on {len(questions_data)} questions...")
+    for item in questions_data:
+        task_id = item.get("task_id")
+        question_text = item.get("question")
+        if not task_id or question_text is None:
+            print(f"Skipping item with missing task_id or question: {item}")
+            continue
+        try:
+            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})
+        except Exception as e:
+             print(f"Error running agent on task {task_id}: {e}")
+             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
 
-        # Process questions
-        answers_payload = []
-        for item in questions_data:
-            task_id = item.get("task_id")
-            question_text = item.get("question")
-            if not task_id or not question_text:
-                continue
-                
-            try:
-                answer = agent(question_text)
-                answers_payload.append({
-                    "task_id": task_id,
-                    "submitted_answer": answer
-                })
-                results_log.append({
-                    "Task ID": task_id,
-                    "Question": question_text,
-                    "Submitted Answer": answer
-                })
-            except Exception as e:
-                results_log.append({
-                    "Task ID": task_id,
-                    "Question": question_text,
-                    "Submitted Answer": f"AGENT ERROR: {e}"
-                })
+    if not answers_payload:
+        print("Agent did not produce any answers to submit.")
+        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
 
-        # Submit answers
-        submission_data = {
-            "username": username.strip(),
-            "agent_code": agent_code,
-            "answers": answers_payload
-        }
-        
-        response = requests.post(f"{api_url}/submit", json=submission_data, timeout=60)
+    # 4. Prepare Submission 
+    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}'..."
+    print(status_update)
+
+    # 5. Submit
+    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
+    try:
+        response = requests.post(submit_url, json=submission_data, timeout=60)
         response.raise_for_status()
         result_data = response.json()
-        
         final_status = (
-            f"Submission Successful!\nOverall Score: {result_data.get('score', 'N/A')}%\n"
-            f"Correct: {result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')}\n"
-            f"Message: {result_data.get('message', 'No message')}"
+            f"Submission Successful!\n"
+            f"User: {result_data.get('username')}\n"
+            f"Overall Score: {result_data.get('score', 'N/A')}% "
+            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
+            f"Message: {result_data.get('message', 'No message received.')}"
         )
-        return final_status, pd.DataFrame(results_log)
-    
+        print("Submission successful.")
+        results_df = pd.DataFrame(results_log)
+        return final_status, results_df
+    except requests.exceptions.HTTPError as e:
+        error_detail = f"Server responded with status {e.response.status_code}."
+        try:
+            error_json = e.response.json()
+            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
+        except requests.exceptions.JSONDecodeError:
+            error_detail += f" Response: {e.response.text[:500]}"
+        status_message = f"Submission Failed: {error_detail}"
+        print(status_message)
+        results_df = pd.DataFrame(results_log)
+        return status_message, results_df
+    except requests.exceptions.Timeout:
+        status_message = "Submission Failed: The request timed out."
+        print(status_message)
+        results_df = pd.DataFrame(results_log)
+        return status_message, results_df
+    except requests.exceptions.RequestException as e:
+        status_message = f"Submission Failed: Network error - {e}"
+        print(status_message)
+        results_df = pd.DataFrame(results_log)
+        return status_message, results_df
     except Exception as e:
-        return f"Error: {str(e)}", pd.DataFrame(results_log)
+        status_message = f"An unexpected error occurred during submission: {e}"
+        print(status_message)
+        results_df = pd.DataFrame(results_log)
+        return status_message, results_df
+
 
-# Gradio UI Setup
+# --- Build Gradio Interface using Blocks ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Basic Agent Evaluation Runner")
+    gr.Markdown("# Advanced Agent Evaluation Runner")
     gr.Markdown(
         """
         **Instructions:**
-        1. Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
-        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.  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.
         ---
-        **Disclaimers:**
-        Once clicking on the "submit button, it can take quite some time (this is the time for the agent to go through all the questions).
-        This space provides a basic setup and is intentionally sub-optimal to encourage you to develop your own, more robust solution. For instance, for the delay process of the submit button, a solution could be to cache the answers and submit in a separate action or even to answer the questions in async.
+        **Note:** 
+        The evaluation might take some time as the agent processes all questions through the tools.
         """
     )
 
@@ -694,16 +287,16 @@ if __name__ == "__main__":
         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?).")
+        print("ℹ️  SPACE_HOST environment variable not found (running locally?).")
 
     if space_id_startup: # Print repo URLs if SPACE_ID is found
         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")
+        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("ℹ️  SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")
 
     print("-"*(60 + len(" App Starting ")) + "\n")
 
-    print("Launching Gradio Interface for Basic Agent Evaluation...")
+    print("Launching Gradio Interface for Advanced Agent Evaluation...")
     demo.launch(debug=True, share=False)