rubenjpdev/gaia-final-assignment-agent

1import itertools2import os3import time4 5from langchain_community.tools import WikipediaQueryRun6from langchain_community.utilities import WikipediaAPIWrapper7from langchain_core.tools import tool8from langchain_ollama import ChatOllama9from langgraph.prebuilt import create_react_agent10from ddgs import DDGS11 12_BACKENDS = itertools.cycle(["brave", "lite", "html", "bing"])13 14 15@tool16def web_search(query: str) -> str:17    """Search the web and return short snippets of the top results."""18    last_err = None19    for attempt in range(4):20        backend = next(_BACKENDS)21        try:22            results = list(DDGS().text(query, max_results=4, backend=backend))23            if results:24                return "\n\n".join(f"{r['title']}: {r['body']}" for r in results)25            last_err = "no results"26        except Exception as e:27            last_err = str(e)28        time.sleep(1.5 * (attempt + 1))29    return f"search failed: {last_err}"30 31MODEL_ID = os.environ.get("AGENT_MODEL", "qwen2:7b")32 33SYSTEM_PROMPT = (34    "You are a general AI assistant answering GAIA benchmark questions. "35    "Use tools when you need facts, math, web content, or text/string manipulation "36    "(e.g. reversing or decoding text) you can't do reliably in your head — use "37    "python_eval for those instead of reasoning it out yourself. "38    "If the question specifies an output format (units, order, comma-separated, "39    "alphabetical, etc.), follow it exactly — that overrides brevity. "40    "Report your reasoning, then finish your last line with exactly:\n"41    "FINAL ANSWER: <answer>\n"42    "Otherwise the answer must be as short as possible: a number, a name, or a "43    "few words, no units unless asked, no explanation after that line."44)45 46 47@tool48def python_eval(expression: str) -> str:49    """Evaluate a Python expression and return the result. Use for arithmetic, counting,50    date math, and string manipulation like expression[::-1] to reverse text."""51    try:52        return str(eval(expression, {"__builtins__": {}}))53    except Exception as e:54        return f"error: {e}"55 56 57def build_agent():58    llm = ChatOllama(model=MODEL_ID, temperature=0)59    tools = [web_search, WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()), python_eval]60    return llm, create_react_agent(llm, tools, prompt=SYSTEM_PROMPT)61 62 63def _extract(text: str) -> str:64    if "FINAL ANSWER:" in text:65        text = text.split("FINAL ANSWER:")[-1]66    return text.strip().strip('"').removeprefix("<nil>").strip()67 68 69_SQUEEZE_PROMPT = (70    "Extract only the final answer from this text, as few words/numbers as "71    "possible, no punctuation, no explanation, no sentence. If it's a number, "72    "give just the number. If it's a name, give just the name.\n\nTEXT:\n{text}"73)74 75 76def answer(agent, question: str) -> str:77    llm, graph = agent78    result = graph.invoke({"messages": [("user", question)]})79    raw = _extract(result["messages"][-1].content)80    if not raw or len(raw.split()) > 8:81        squeezed = llm.invoke(_SQUEEZE_PROMPT.format(text=raw or result["messages"][-1].content))82        raw = squeezed.content.strip().strip('"')83    return raw84 85 86if __name__ == "__main__":87    a = build_agent()88    print(answer(a, "What is 12 * 7 + 3? Explain then give FINAL ANSWER: <the number>"))89