Update app.py

app.py

@@ -4,198 +4,32 @@ import requests
 import pandas as pd
 from langchain.agents import create_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_openai import ChatOpenAI
-from langchain.tools import tool
+# Agent implementation is moved to gaia_agent.py
+from gaia_agent import GAIAAgent
 from dotenv import load_dotenv
-from langchain_community.document_loaders import ArxivLoader, WikipediaLoader
-from ddgs import DDGS
 
 # Load environment variables
-# load_dotenv()
+load_dotenv()
 
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Agent Setup ---
-#openai_key = os.getenv("OPENAI_API_KEY")
+openai_key = os.getenv("OPENAI_API_KEY")
 googleai_key = os.getenv("GOOGLE_API_KEY")
 
-# Initialize the model
-model = ChatGoogleGenerativeAI(
-    model="gemini-2.5-flash",
-    temperature=0,
-    max_tokens=5000,
-    timeout=None,
-    max_retries=2,
-)
+# Use OpenRouter via LangChain's ChatOpenAI
+openrouter_key = os.getenv("OPENROUTER_API_KEY")
+if not openrouter_key:
+    raise RuntimeError("Set OPENROUTER_API_KEY in your .env (OpenRouter API key)")
 
-# --- Tools Definition ---
-@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}
+# model is created inside gaia_agent module
+# (gaia_agent.py will initialize the ChatOpenAI model using OPENROUTER_API_KEY)
 
-@tool
-def web_search(query: str) -> str:
-    """Search DDGS for a query and return maximum 3 results.
-    
-    Args:
-        query: The search query."""
-    search_docs = DDGS().text(query,max_results=3)
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'Title:{doc["title"]}\nContent:{doc["body"]}\n--\n'
-            for doc in search_docs
-        ])
-    return 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}
-
-@tool
-def image_search(query: str) -> str:
-    """Searches DDGS for an image query and returns maximum 10 image results"""
-    search_images = DDGS().images(query=query)
-    formatted_result = "\n\n---\n\n".join(
-        [
-            f'Image Title:{image["title"]}\nImage URL: {image["url"]}'
-            for image in search_images
-        ])
 
 
-# Tools list
-tools = [
-    multiply, add, subtract, divide, modulus,
-    wiki_search, web_search, arvix_search, image_search
-]
-
-# System prompt
-sys_prompt = """You are a helpful agent, please provide clear and concise answers to asked questions.
-Keep your word limit for answers as minimum as you can. You are equipped with the following tools:
-1. [multiply], [add], [subtract], [divide], [modulus] - basic calculator operations.
-2. [wiki_search] - search Wikipedia and return up to 2 documents as text.
-3. [web_search] - perform a web search and return up to 3 documents as text.
-4. [arxiv_search] - search arXiv and return up to 3 documents as text.
-5. [image_search] - Searches the internet for an image query and returns maximum 10 image results
-
-Under any circumstances, if you fail to provide the accurate answer expected by the user, you may say the same to the user and provide a similar answer which is approximately the closest. Disregard spelling mistakes and provide answer with results retreived from the correct spelling. 
-
-For every tool you use, append a single line at the end of your response exactly in this format:
-[TOOLS USED: (tool_name)]
-When no tools are used, append:
-[TOOLS USED WERE NONE]"""
-
-# --- Agent Class ---
-class GAIAAgent:
-    def __init__(self):
-        print("GAIAAgent initialized with LangChain agent.")
-        try:
-            self.agent = create_agent(model, tools=tools, system_prompt=sys_prompt)
-            print("Agent created successfully.")
-        except Exception as e:
-            print(f"Error creating agent: {e}")
-            raise
-    
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 100 chars): {question[:100]}...")
-        try:
-            result = self.agent.invoke({
-                "messages": [{"role": "user", "content": question}]
-            })
-            
-            # Get the content from the last message
-            raw_content = result["messages"][-1].content
-            
-            # Parse the response format: list of dicts with 'text' key
-            if isinstance(raw_content, list) and len(raw_content) > 0:
-                if isinstance(raw_content[0], dict) and 'text' in raw_content[0]:
-                    answer = raw_content[0]['text']
-                else:
-                    # Fallback: convert list to string
-                    answer = str(raw_content)
-            elif isinstance(raw_content, str):
-                answer = raw_content
-            else:
-                answer = str(raw_content)
-            
-            print(f"Agent returning answer (first 100 chars): {answer[:100]}...")
-            return answer
-        except Exception as e:
-            print(f"Error in agent execution: {e}")
-            import traceback
-            traceback.print_exc()
-            return f"Error: {str(e)}"
+# The tools and GAIAAgent implementation live in gaia_agent.py now. This file
+# imports GAIAAgent and uses it in run_and_submit_all.
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     """
@@ -372,6 +206,7 @@ if __name__ == "__main__":
     space_host_startup = os.getenv("SPACE_HOST")
     space_id_startup = os.getenv("SPACE_ID")
     google_api_key = os.getenv("GOOGLE_API_KEY")
+    tavily_api_key = os.getenv("TAVILY_API_KEY")
 
     if space_host_startup:
         print(f"✅ SPACE_HOST found: {space_host_startup}")
@@ -390,6 +225,7 @@ if __name__ == "__main__":
         print("✅ GOOGLE_API_KEY found")
     else:
         print("⚠️  GOOGLE_API_KEY not found - agent will not work without it!")
+
     print("-"*(60 + len(" App Starting ")) + "\n")
 
     print("Launching Gradio Interface for GAIA Agent Evaluation...")