agent.py

# agent.py — updated with ReActAgent (new API) and improved system prompt

import asyncio
from llama_index.llms.openai import OpenAI
from llama_index.core.tools import FunctionTool
from llama_index.core.agent.react.base import ReActAgent

from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
from langchain_experimental.tools.python.tool import PythonREPLTool
from langchain_community.document_loaders import YoutubeLoader

# --- Tool Wrappers ---
def search_duckduckgo(query: str) -> str:
    return DuckDuckGoSearchRun().run(query)

def search_wikipedia(query: str) -> str:
    return WikipediaQueryRun(api_wrapper=None).run(query)

def run_python(code: str) -> str:
    return PythonREPLTool().run(code)

def get_youtube_transcript(url: str) -> str:
    loader = YoutubeLoader.from_youtube_url(url, add_video_info=False)
    docs = loader.load()
    return " ".join(doc.page_content for doc in docs)

# --- Define tools ---
TOOLS = [
    FunctionTool.from_defaults(search_duckduckgo),
    FunctionTool.from_defaults(search_wikipedia),
    FunctionTool.from_defaults(run_python),
    FunctionTool.from_defaults(get_youtube_transcript),
]

# --- Load LLM ---
llm = OpenAI(model="gpt-4")

# --- Create ReAct Agent (new API) ---
agent = ReActAgent.from_tools(
    tools=TOOLS,
    llm=llm,
    verbose=False,
    system_prompt="""
You are an expert AI assistant participating in the GAIA benchmark.

Your task is to answer questions as precisely, concisely, and correctly as possible. Each answer will be evaluated by an automated system that checks both correctness and strict formatting.

Follow these rules:

1. Focus only on the final answer — no explanation, no reasoning, no commentary.
2. Format your answer exactly as requested:
   - If the answer must be a list, output a **comma-separated list** in alphabetical order with no extra words.
   - If the answer is a **chess move**, return it in **algebraic notation** like `Qd1+` or `Rf3#` with no spaces or commentary.
   - If the answer is a **number**, return **only the number** (e.g. `5`, `33`, `0`), not `five` or `The answer is 5`.
   - If the answer is a **currency value**, return it as a USD amount formatted like `$10.25`.
   - If the answer is a **name**, return only the first name or surname as required, with no quotes or additional text.
   - If the answer must be from **an Excel file or audio or image**, analyze the file or transcript and return only the required value.

3. If the question is ambiguous, give your best guess using available tools. Do not leave answers blank.

You have access to tools such as:
- internet search (DuckDuckGo)
- Wikipedia lookup
- Python code interpreter
- YouTube transcript parser

Always choose the tool(s) most relevant to the question.

Be accurate, be concise, and never guess the format — follow it exactly.
"""
)

# --- Runner with timeout ---
async def answer_question(question: str) -> str:
    try:
        result = await asyncio.wait_for(agent.aget_response(question), timeout=60)
        return str(result)
    except asyncio.TimeoutError:
        return "[TIMEOUT]"
    except Exception as e:
        return f"[ERROR] {e}"