Update agent.py

agent.py

@@ -1,188 +1,219 @@
+# --- Basic Agent Definition ---
+import asyncio
 import os
-from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import ToolNode, tools_condition
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_groq import ChatGroq
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
-from langchain_core.messages import SystemMessage, HumanMessage
-from langchain_core.tools import tool
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_community.document_loaders import WikipediaLoader
-from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound
-from duckduckgo_search import DDGS
-from langchain_community.document_loaders import ArxivLoader
-from sympy import sympify
-from PIL import Image
-import re
+import sys
+import logging
+import random
+import pandas as pd
 import requests
-from io import BytesIO
+import wikipedia as wiki
+from markdownify import markdownify as to_markdown
+from typing import Any
 from dotenv import load_dotenv
+from google.generativeai import types, configure
 
+from smolagents import InferenceClientModel, LiteLLMModel, ToolCallingAgent, Tool, DuckDuckGoSearchTool
+
+# Load environment and configure Gemini
 load_dotenv()
+configure(api_key=os.getenv("GOOGLE_API_KEY"))
+
+# Logging
+#logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
+#logger = logging.getLogger(__name__)
+
+# --- Model Configuration ---
+GEMINI_MODEL_NAME = "gemini/gemini-1.5-flash"
+OPENAI_MODEL_NAME = "openai/gpt-4o"
+GROQ_MODEL_NAME = "groq/llama3-70b-8192"
+DEEPSEEK_MODEL_NAME = "deepseek/deepseek-chat"
+HF_MODEL_NAME = "Qwen/Qwen2.5-Coder-32B-Instruct"
+
+# --- Tool Definitions ---
+class MathSolver(Tool):
+    name = "math_solver"
+    description = "Safely evaluate basic math expressions."
+    inputs = {"input": {"type": "string", "description": "Math expression to evaluate."}}
+    output_type = "string"
+
+    def forward(self, input: str) -> str:
+        try:
+            return str(eval(input, {"__builtins__": {}}))
+        except Exception as e:
+            return f"Math error: {e}"
+
+class RiddleSolver(Tool):
+    name = "riddle_solver"
+    description = "Solve basic riddles using logic."
+    inputs = {"input": {"type": "string", "description": "Riddle prompt."}}
+    output_type = "string"
+
+    def forward(self, input: str) -> str:
+        if "forward" in input and "backward" in input:
+            return "A palindrome"
+        return "RiddleSolver failed."
+
+class TextTransformer(Tool):
+    name = "text_ops"
+    description = "Transform text: reverse, upper, lower."
+    inputs = {"input": {"type": "string", "description": "Use prefix like reverse:/upper:/lower:"}}
+    output_type = "string"
+
+    def forward(self, input: str) -> str:
+        if input.startswith("reverse:"):
+            reversed_text = input[8:].strip()[::-1]
+            if 'left' in reversed_text.lower():
+                return "right"
+            return reversed_text
+        if input.startswith("upper:"):
+            return input[6:].strip().upper()
+        if input.startswith("lower:"):
+            return input[6:].strip().lower()
+        return "Unknown transformation."
+
+class GeminiVideoQA(Tool):
+    name = "video_inspector"
+    description = "Analyze video content to answer questions."
+    inputs = {
+        "video_url": {"type": "string", "description": "URL of video."},
+        "user_query": {"type": "string", "description": "Question about video."}
+    }
+    output_type = "string"
+
+    def __init__(self, model_name, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model_name = model_name
+
+    def forward(self, video_url: str, user_query: str) -> str:
+        req = {
+            'model': f'models/{self.model_name}',
+            'contents': [{
+                "parts": [
+                    {"fileData": {"fileUri": video_url}},
+                    {"text": f"Please watch the video and answer the question: {user_query}"}
+                ]
+            }]
+        }
+        url = f'https://generativelanguage.googleapis.com/v1beta/models/{self.model_name}:generateContent?key={os.getenv("GOOGLE_API_KEY")}'
+        res = requests.post(url, json=req, headers={'Content-Type': 'application/json'})
+        if res.status_code != 200:
+            return f"Video error {res.status_code}: {res.text}"
+        parts = res.json()['candidates'][0]['content']['parts']
+        return "".join([p.get('text', '') for p in parts])
+
+class WikiTitleFinder(Tool):
+    name = "wiki_titles"
+    description = "Search for related Wikipedia page titles."
+    inputs = {"query": {"type": "string", "description": "Search query."}}
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        results = wiki.search(query)
+        return ", ".join(results) if results else "No results."
+
+class WikiContentFetcher(Tool):
+    name = "wiki_page"
+    description = "Fetch Wikipedia page content."
+    inputs = {"page_title": {"type": "string", "description": "Wikipedia page title."}}
+    output_type = "string"
+
+    def forward(self, page_title: str) -> str:
+        try:
+            return to_markdown(wiki.page(page_title).html())
+        except wiki.exceptions.PageError:
+            return f"'{page_title}' not found."
 
-# Load system prompt
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    SYSTEM_PROMPT = f.read()
-
-# Tool: Wikipedia search
-@tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return content from up to 2 documents."""
-    try:
-        docs = WikipediaLoader(query=query, load_max_docs=2).load()
-        return "\n\n---\n\n".join([doc.page_content for doc in docs])
-    except Exception as e:
-        return f"Wikipedia search failed: {e}"
-
-# Tool: Tavily web search
-@tool
-def web_search(query: str) -> str:
-    """Search the web using Tavily and return content from up to 3 results."""
-    try:
-        results = TavilySearchResults(max_results=3).invoke(query)
-        if isinstance(results, list):
-            return "\n\n---\n\n".join([r["content"] if isinstance(r, dict) else str(r) for r in results])
-        return str(results)
-    except Exception as e:
-        return f"Web search failed: {e}"
-
-# Tool: DuckDuckGo search
-@tool
-def duckduckgo_search(query: str) -> str:
-    """Search using DuckDuckGo and return summaries from up to 3 results."""
-    try:
-        with DDGS() as ddgs:
-            results = ddgs.text(query, max_results=3)
-            return "\n\n---\n\n".join([r["body"] for r in results if "body" in r])
-    except Exception as e:
-        return f"DuckDuckGo search failed: {e}"
-
-# Tool: YouTube transcript or duration extractor
-@tool
-def youtube_transcript(video_title_or_url: str) -> str:
-    """Get duration of a YouTube video using its title or URL."""
-    try:
-        with DDGS() as ddgs:
-            results = ddgs.videos(video_title_or_url, max_results=1)
-            if not results:
-                return "No video found by that title."
-            video = results[0]
-            return f"Duration: {video.get('duration')}"
-    except Exception as e:
-        return f"YouTube search failed: {e}"
-
-# Tool: Arxiv paper fetcher
-@tool
-def arxiv_fetch(query_or_id: str) -> str:
-    """Fetch metadata from arXiv either by ID or search query."""
-    try:
-        if re.match(r"\d{4}\.\d{5}(v\d+)?", query_or_id):
-            abs_url = f"https://arxiv.org/abs/{query_or_id}"
-            api_url = f"http://export.arxiv.org/api/query?id_list={query_or_id}"
-            res = requests.get(api_url)
-            if res.status_code == 200:
-                return res.text[:2000] + f"\n\nFull: {abs_url}"
-            return f"Could not retrieve metadata from arXiv API"
-        else:
-            docs = ArxivLoader(query=query_or_id, load_max_docs=2).load()
-            return "\n\n---\n\n".join([doc.page_content for doc in docs])
-    except Exception as e:
-        return f"ArXiv fetch failed: {e}"
-
-@tool
-def math_solver(expression: str) -> str:
-    """Evaluate a math expression and return the result."""
-    try:
-        result = sympify(expression).evalf()
-        return str(result)
-    except Exception as e:
-        return f"Math error: {e}"
-
-@tool
-def reverse_text(text: str) -> str:
-    """Reverse the input string."""
-    return text[::-1]
-
-@tool
-def image_info(url: str) -> str:
-    """Fetch image size (width x height) from a given URL."""
-    try:
-        response = requests.get(url)
-        img = Image.open(BytesIO(response.content))
-        return f"Image size: {img.size} (width x height)"
-    except Exception as e:
-        return f"Image error: {e}"
-
-# Tools list
-tools = [
-    wiki_search,
-    web_search,
-    duckduckgo_search,
-    youtube_transcript,
-    arxiv_fetch,
-    math_solver,
-    reverse_text,
-    image_info
-]
-
-def build_graph(provider: str = "groq"):
-    if provider == "google":
-        llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
-    elif provider == "groq":
-        llm = ChatGroq(model="llama3-70b-8192", temperature=0)
-    elif provider == "huggingface":
-        llm = ChatHuggingFace(
-            llm=HuggingFaceEndpoint(
-                url="https://api-inference.huggingface.co/models/tiiuae/falcon-7b-instruct",
-                temperature=0,
-            ),
+# --- Basic Agent Definition ---
+class BasicAgent:
+    def __init__(self, provider="deepseek"):
+        print("BasicAgent initialized.")
+        model = self.select_model(provider)
+        client = InferenceClientModel()
+        tools = [
+            DuckDuckGoSearchTool(),
+            GeminiVideoQA(GEMINI_MODEL_NAME),
+            WikiTitleFinder(),
+            WikiContentFetcher(),
+            MathSolver(),
+            RiddleSolver(),
+            TextTransformer(),
+        ]
+        self.agent = ToolCallingAgent(
+            model=model,
+            tools=tools,
+            add_base_tools=False,
+            max_steps=5,
+        )
+        self.agent.system_prompt = (
+            """
+            You are a general AI assistant. I will ask you a question. 
+            YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
+            If your answer is a number and you are not explicitly asked for a string, write it in numerals instead of words, and don't use comma to write your number nor 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.
+
+            Answer questions as literally as you can, making as few assumptions as possible. Restrict the answer to the narrowest definition that still satifies the question.
+            If you are provied with a video, please watch and summarize the entire video before answering the question. The correct answer may be present only in a few frames of the video.
+            If you have difficulty finding an answer on Wikipedia, you may search the internet using Google Search or Duckduckgo search.
+            If you are asked to prove something, first state your assumptions and think step by step before giving your final answer.
+
+            Your final answer must strictly follow this format:
+            FINAL ANSWER: [ANSWER]
+
+            Only write the answer in that exact format. Do not explain anything. Do not include any other text.
+            """
         )
-    else:
-        raise ValueError("Invalid provider. Choose 'google', 'groq', or 'huggingface'.")
-
-    llm_with_tools = llm.bind_tools(tools)
-
-    def system_node(state: MessagesState):
-        return {"messages": [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]}
-
-    def assistant_node(state: MessagesState):
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
-
-    builder = StateGraph(MessagesState)
-    builder.add_node("system", system_node)
-    builder.add_node("assistant", assistant_node)
-    builder.add_node("tools", ToolNode(tools))
-    builder.add_edge(START, "system")
-    builder.add_edge("system", "assistant")
-    builder.add_conditional_edges("assistant", tools_condition)
-    builder.add_edge("tools", "assistant")
-    return builder.compile()
 
-class BasicAgent:
-    def __init__(self, provider="groq"):
-        print(f"GAIA LangGraph Agent Initialized using {provider}")
-        self.graph = build_graph(provider)
+    def select_model(self, provider: str):
+        if provider == "openai":
+            return LiteLLMModel(model_id=OPENAI_MODEL_NAME, api_key=os.getenv("OPENAI_API_KEY"))
+        elif provider == "groq":
+            return LiteLLMModel(model_id=GROQ_MODEL_NAME, api_key=os.getenv("GROQ_API_KEY"))
+        elif provider == "deepseek":
+            return LiteLLMModel(model_id=DEEPSEEK_MODEL_NAME, api_key=os.getenv("DEEPSEEK_API_KEY"))
+        elif provider == "hf":
+            return InferenceClientModel()
+        else:
+            return LiteLLMModel(model_id=GEMINI_MODEL_NAME, api_key=os.getenv("GOOGLE_API_KEY"))
 
     def __call__(self, question: str) -> str:
-        try:
-            messages = [HumanMessage(content=question)]
-            result = self.graph.invoke({"messages": messages})
-            final_msg = result["messages"][-1].content.strip()
-            if not final_msg.startswith("FINAL ANSWER:"):
-                final_msg = f"FINAL ANSWER: {final_msg}"
-            return final_msg
-        except Exception as e:
-            return f"FINAL ANSWER: error - {str(e)}"
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        result = self.agent.run(question)
+        if isinstance(result, dict) and "final_answer" in result and isinstance(result["final_answer"], str):
+            final_str = result["final_answer"].strip()
+        else:
+            final_str = str(result).strip()
+            
+        return final_str
+
+    def evaluate_random_questions(self, csv_path: str = "gaia_qa.csv", sample_size: int = 3):
+        df = pd.read_csv(csv_path)
+        if not {"question", "answer"}.issubset(df.columns):
+            print("CSV must contain 'question' and 'answer' columns.")
+            print("Found columns:", df.columns.tolist())
+            return
+        samples = df.sample(n=sample_size)
+        for _, row in samples.iterrows():
+            question = row["question"].strip()
+            expected = f"FINAL ANSWER: {str(row['answer']).strip()}"
+            result = self(question).strip()
+            print("---")
+            print("Question:", question)
+            print("Expected:", expected)
+            print("Agent:", result)
+            print("Correct:", expected == result)
 
 if __name__ == "__main__":
-    agent = BasicAgent(provider="groq")
-    questions = [
-        "What is the zip code of the Eiffel Tower?",
-        "What is the capital city of Australia?",
-        "How long is the video titled 'The History of Time' on YouTube?",
-        "What does the arXiv paper '2303.12712' say about Transformer performance?",
-    ]
-
-    for q in questions:
-        print(f"\n[Question]: {q}")
+    args = sys.argv[1:]
+    if not args or args[0] in {"-h", "--help"}:
+        print("Usage: python agent.py [question | dev]\n")
+        print(" - Provide a question to get a GAIA-style answer.")
+        print(" - Use 'dev' to evaluate 3 random GAIA questions from gaia_qa.csv.")
+        sys.exit(0)
+
+    q = " ".join(args)
+    agent = BasicAgent()
+    if q == "dev":
+        agent.evaluate_random_questions()
+    else:
         print(agent(q))