Upload gaia_agent.py

gaia_agent.py

@@ -0,0 +1,161 @@
+import os
+import requests
+from langchain.agents import create_agent
+from langchain.tools import tool
+from dotenv import load_dotenv
+from langchain_community.document_loaders import ArxivLoader, WikipediaLoader
+from ddgs import DDGS
+from bs4 import BeautifulSoup
+
+# Load environment variables
+load_dotenv()
+
+# --- Agent Setup ---
+openai_key = os.getenv("OPENAI_API_KEY")
+googleai_key = os.getenv("GOOGLE_API_KEY")
+
+# 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)")
+
+# Defer ChatOpenAI import until runtime to avoid import-time errors in environments without the package
+from langchain_openai import ChatOpenAI
+
+model = ChatOpenAI(
+  api_key=openrouter_key,
+  base_url="https://openrouter.ai/api/v1",
+  model="gpt-4o-mini",
+  max_completion_tokens=10000,
+)
+
+# --- Tools Definition ---
+@tool
+def multiply(a: int, b: int) -> int:
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> float:
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    return a % b
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results."""
+    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 formatted_search_docs
+
+@tool
+def web_search(query: str) -> str:
+    """Search DDGS for a query and return maximum 3 results."""
+    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 arxiv_search(query: str) -> str:
+    """Search arXiv for a query and return maximum 3 results."""
+    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 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
+        ]
+    )
+    return formatted_result
+
+@tool
+def fetch_url_content(url: str) -> str:
+    """Fetch and return the text content from a webpage URL."""
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        soup = BeautifulSoup(response.text, 'html.parser')
+        for script in soup(["script", "style"]):
+            script.decompose()
+        text = soup.get_text(separator='\n', strip=True)
+        return text[:2000] + ("..." if len(text) > 2000 else "")
+    except Exception as e:
+        return f"Error fetching URL: {str(e)}"
+
+# Tools list
+tools = [
+    multiply, add, subtract, divide, modulus,
+    wiki_search, web_search, arxiv_search, image_search,
+    fetch_url_content,
+]
+
+# 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]
+"""
+
+class GAIAAgent:
+    def __init__(self):
+        # create internal agent
+        try:
+            self.agent = create_agent(model, tools=tools, system_prompt=sys_prompt)
+        except Exception as e:
+            raise
+
+    def __call__(self, question: str) -> str:
+        result = self.agent.invoke({"messages": [{"role": "user", "content": question}]})
+        raw_content = result["messages"][-1].content
+        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:
+                answer = str(raw_content)
+        elif isinstance(raw_content, str):
+            answer = raw_content
+        else:
+            answer = str(raw_content)
+        return answer