add agent

agent.py

@@ -0,0 +1,359 @@
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain.tools import tool
+
+from langchain_community.document_loaders import WikipediaLoader,ArxivLoader
+
+from tavily import TavilyClient
+
+from openai import OpenAI
+import base64
+import re
+import os
+
+
+from typing import TypedDict, Annotated, Literal
+ 
+from langchain_core.messages import (
+    AnyMessage, HumanMessage, AIMessage, ToolMessage, SystemMessage
+)
+
+from langgraph.graph.message import add_messages
+from langgraph.graph import StateGraph, END
+
+
+
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+TAVILY_API_KEY = os.environ.get("TAVILY_API_KEY")
+
+
+tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
+openai_client = OpenAI(api_key=OPENAI_API_KEY)
+
+MAX_STEPS = 15
+
+@tool
+def search_wikipedia(query: str, max_docs: int = 3) -> str:
+    """Search Wikipedia for general knowledge and return summarized content.
+ 
+    Args:
+        query: Topic to search (e.g., 'Artificial Intelligence', 'France history')
+        max_docs: Maximum number of Wikipedia pages to retrieve
+    """
+    loader = WikipediaLoader(query=query, load_max_docs=max_docs)
+    docs = loader.load()
+    return "\n\n".join(doc.page_content[:3000] for doc in docs)
+
+
+@tool
+def search_arxiv(query: str, max_docs: int = 3) -> str:
+    """Search arXiv for scientific papers and return summaries.
+ 
+    Args:
+        query: Research topic or keywords (e.g., 'transformer attention')
+        max_docs: Maximum number of papers to retrieve
+    """
+    loader = ArxivLoader(query=query, load_max_docs=max_docs)
+    docs = loader.load()
+    return "\n\n".join(doc.page_content[:3000] for doc in docs)
+
+
+@tool
+def search_web(query: str, max_results: int = 5) -> str:
+    """Search the web for up-to-date information.
+ 
+    Args:
+        query: Search query (e.g., 'latest OpenAI model 2025')
+        max_results: Number of results to return
+    """
+    response = tavily_client.search(query=query, max_results=max_results)
+    results = [f"{r['title']}\n{r['content']}" for r in response["results"]]
+    return "\n\n".join(results)
+
+
+
+@tool
+def transcribe_audio(file_path: str) -> str:
+    """Transcribe an audio file (mp3, wav) into text.
+ 
+    Args:
+        file_path: Path to the audio file on disk
+    """
+    with open(file_path, "rb") as f:
+        transcript = openai_client.audio.transcriptions.create(
+            model="whisper-1",
+            file=f,
+        )
+    return transcript.text
+
+
+@tool
+def read_image(file_path: str) -> str:
+    """Read an image file and return a description via GPT-4o vision.
+ 
+    Args:
+        file_path: Path to the image file on disk
+    """
+    with open(file_path, "rb") as f:
+        b64 = base64.b64encode(f.read()).decode("utf-8")
+    ext = file_path.rsplit(".", 1)[-1].lower()
+    mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg",
+            "png": "image/png", "gif": "image/gif",
+            "webp": "image/webp"}.get(ext, "image/png")
+    response = openai_client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image_url",
+                     "image_url": {"url": f"data:{mime};base64,{b64}"}},
+                    {"type": "text",
+                     "text": "Describe this image in detail. Extract any text, data, or key information visible."},
+                ],
+            }
+        ],
+        max_tokens=1024,
+    )
+    return response.choices[0].message.content
+
+@tool
+def read_file(file_path: str) -> str:
+    """Read a file and return its contents."""
+    with open(file_path, "r", encoding="utf-8") as f:
+        return f.read()
+
+@tool
+def python_repl(code: str) -> str:
+    """Execute Python code and return stdout + the value of the last expression.
+    Useful for arithmetic, data manipulation, and logic tasks.
+ 
+    Args:
+        code: Valid Python code string
+    """
+    import io, sys, traceback
+    stdout_capture = io.StringIO()
+    local_vars: dict = {}
+    try:
+        sys.stdout = stdout_capture
+        exec(code, {}, local_vars)           # run all lines
+        # try to eval last line as expression
+        lines = [l for l in code.strip().splitlines() if l.strip()]
+        last_val = ""
+        if lines:
+            try:
+                last_val = repr(eval(lines[-1], {}, local_vars))
+            except Exception:
+                pass
+    except Exception:
+        return traceback.format_exc()
+    finally:
+        sys.stdout = sys.__stdout__
+    out = stdout_capture.getvalue()
+    return "\n".join(filter(None, [out, last_val])) or "Code executed successfully (no output)."
+
+
+
+
+TOOLS = [
+    search_wikipedia,
+    search_arxiv,
+    search_web,
+    transcribe_audio,
+    read_image,
+    read_file,
+    python_repl,
+]
+ 
+TOOL_MAP = {t.name: t for t in TOOLS}
+
+
+SYSTEM_PROMPT = f"""You are a highly capable AI assistant solving tasks from the GAIA benchmark.
+ 
+## Core rules (MUST follow)
+1. THINK before acting: decompose the question and plan which tool(s) you need.
+2. NEVER call the same tool with the exact same arguments twice.
+   If the result was insufficient, use a DIFFERENT query or a DIFFERENT tool.
+3. If search_wikipedia returns a biography page instead of a discography/list,
+   immediately switch to search_web with a more specific query.
+4. For calculations / counting, always use python_repl — never guess numbers.
+5. Once you have enough information, STOP calling tools and give the final answer.
+6. You have at most {MAX_STEPS} tool-call rounds total. Budget them wisely.
+ 
+## Tool selection guide
+- General facts / biography  → search_wikipedia (vary query if first try fails)
+- Discographies, filmographies, lists → search_web (Wikipedia tool may miss these)
+- Current events / live data → search_web
+- Scientific papers          → search_arxiv
+- Arithmetic / logic         → python_repl
+- Provided image file        → read_image
+- Provided audio file        → transcribe_audio
+- Provided text/csv/json     → read_file
+ 
+## Answer format
+End your FINAL response with exactly:
+FINAL ANSWER: <your answer>
+ 
+Keep it concise — no units unless asked, lists comma-separated.
+"""
+
+
+class AgentState(TypedDict):
+    messages:   Annotated[list[AnyMessage], add_messages]
+    step_count: int   # counts agent_node invocations
+
+
+def make_llm(model: str = "gpt-5.4-mini") -> ChatOpenAI:
+    return ChatOpenAI(
+        model=model,
+        temperature=0,
+        api_key=OPENAI_API_KEY,
+    ).bind_tools(TOOLS)
+ 
+ 
+llm_with_tools = make_llm()
+
+_step = 0  # console display counter
+ 
+CYAN   = "\033[96m"
+GREEN  = "\033[92m"
+YELLOW = "\033[93m"
+RED    = "\033[91m"
+BOLD   = "\033[1m"
+RESET  = "\033[0m"
+ 
+ 
+def _log(label: str, text: str, color: str = RESET) -> None:
+    print(f"{color}{'─'*60}{RESET}")
+    print(f"{color}[Step {_step}] {label}{RESET}")
+    if text.strip():
+        print(f"{color}{text.strip()}{RESET}")
+
+
+def agent_node(state: AgentState) -> AgentState:
+    global _step
+    _step += 1
+    step_count = state.get("step_count", 0) + 1
+ 
+    messages = state["messages"]
+ 
+    # Inject system prompt on first turn
+    if not any(isinstance(m, SystemMessage) for m in messages):
+        messages = [SystemMessage(content=SYSTEM_PROMPT)] + messages
+ 
+    # Warn model to wrap up when approaching the limit
+    if step_count >= MAX_STEPS - 2:
+        messages = list(messages) + [HumanMessage(
+            content=f"⚠️ You have used {step_count}/{MAX_STEPS} steps. "
+                    "Do NOT call any more tools. Synthesise what you have and give FINAL ANSWER now."
+        )]
+ 
+    _log("🤖 AGENT THINKING …", "", CYAN)
+    response = llm_with_tools.invoke(messages)
+ 
+    if response.content:
+        _log("🤖 AGENT RESPONSE", str(response.content)[:600], CYAN)
+ 
+    if response.tool_calls:
+        calls_summary = "\n".join(
+            f"  • {tc['name']}({', '.join(f'{k}={repr(v)}' for k, v in tc['args'].items())})"
+            for tc in response.tool_calls
+        )
+        _log("🔧 TOOL CALLS PLANNED", calls_summary, YELLOW)
+    else:
+        _log("✅ AGENT FINISHED (no more tool calls)", "", GREEN)
+ 
+    return {"messages": [response], "step_count": step_count}
+ 
+
+
+def tool_node(state: AgentState) -> AgentState:
+    global _step
+    last_msg: AIMessage = state["messages"][-1]
+    tool_results: list[ToolMessage] = []
+ 
+    for tc in last_msg.tool_calls:
+        _step += 1
+        tool_fn = TOOL_MAP.get(tc["name"])
+        _log(f"⚙️  RUNNING: {tc['name']}",
+             "\n".join(f"  {k}: {repr(v)}" for k, v in tc["args"].items()),
+             YELLOW)
+ 
+        if tool_fn is None:
+            result = f"ERROR: unknown tool '{tc['name']}'"
+            _log("❌ TOOL ERROR", result, RED)
+        else:
+            try:
+                result = tool_fn.invoke(tc["args"])
+                preview = str(result)[:500] + ("…" if len(str(result)) > 500 else "")
+                _log(f"📥 RESULT: {tc['name']}", preview, GREEN)
+            except Exception as exc:
+                result = f"ERROR calling {tc['name']}: {exc}"
+                _log(f"❌ TOOL ERROR: {tc['name']}", result, RED)
+ 
+        tool_results.append(
+            ToolMessage(content=str(result), tool_call_id=tc["id"])
+        )
+ 
+    return {"messages": tool_results}
+
+
+
+def should_continue(state: AgentState) -> Literal["tools", "end"]:
+    step_count = state.get("step_count", 0)
+ 
+    if step_count >= MAX_STEPS:
+        print(f"{RED}{'─'*60}")
+        print(f"⛔ MAX_STEPS ({MAX_STEPS}) reached — forcing end.{RESET}")
+        return "end"
+ 
+    last = state["messages"][-1]
+    if isinstance(last, AIMessage) and last.tool_calls:
+        return "tools"
+ 
+    return "end"
+
+
+def build_graph() -> StateGraph:
+    g = StateGraph(AgentState)
+    g.add_node("agent", agent_node)
+    g.add_node("tools", tool_node)
+    g.set_entry_point("agent")
+    g.add_conditional_edges("agent", should_continue, {"tools": "tools", "end": END})
+    g.add_edge("tools", "agent")   # always return to agent after tool use
+    return g.compile()
+ 
+ 
+graph = build_graph()
+
+
+def run_agent(question: str, file_path: str | None = None) -> str:
+    """Run the agent on a GAIA question and return the extracted final answer."""
+    global _step
+    _step = 0
+ 
+    print(f"\n{BOLD}{'═'*60}{RESET}")
+    print(f"{BOLD}❓ QUESTION: {question}{RESET}")
+    if file_path:
+        print(f"{BOLD}📎 FILE: {file_path}{RESET}")
+    print(f"{BOLD}{'═'*60}{RESET}\n")
+ 
+    content = question
+    if file_path:
+        content += f"\n\n[Attached file available at: {file_path}]"
+ 
+    result = graph.invoke({
+        "messages": [HumanMessage(content=content)],
+        "step_count": 0,
+    })
+ 
+    last_msg = result["messages"][-1]
+    text = last_msg.content if isinstance(last_msg, AIMessage) else str(last_msg)
+ 
+    match = re.search(r"FINAL ANSWER:\s*(.+)", text, re.IGNORECASE | re.DOTALL)
+    answer = match.group(1).strip() if match else text.strip()
+ 
+    print(f"\n{BOLD}{GREEN}{'═'*60}{RESET}")
+    print(f"{BOLD}{GREEN}🏁 FINAL ANSWER: {answer}{RESET}")
+    print(f"{BOLD}{GREEN}{'═'*60}{RESET}\n")
+    return answer