Update agent.py

agent.py

@@ -1,241 +1,197 @@
-# agent.py  –  Gemini 2.0 Flash · LangGraph · Mehrere Tools
-# =========================================================
-import os, asyncio, base64, mimetypes, tempfile, functools, json
-from typing import Dict, Any, List, Optional
+# agent.py  –  LangChain · LangGraph · Gemini Flash
+# ================================================
+
+"""
+Abhängigkeiten (requirements.txt):
+----------------------------------
+langchain==0.1.*
+langgraph
+google-generativeai
+tavily-python
+wikipedia-api
+pandas
+openpyxl
+tabulate
+"""
+
+import os, re, time, functools
+from typing import Dict, Any, List
 
-from langgraph.graph import START, StateGraph, MessagesState, END
-from langgraph.prebuilt import tools_condition, ToolNode
+import pandas as pd
+from langgraph.graph import StateGraph, START, END, MessagesState
+from langgraph.prebuilt import ToolNode, tools_condition
 
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
-
-from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
+from langchain.tools.python.tool import PythonAstREPLTool
 
 # ---------------------------------------------------------------------
-# Konstanten / API-Keys
+# 0) Optionale LangSmith-Tracing  (setze ENV: LANGCHAIN_API_KEY)
 # ---------------------------------------------------------------------
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
-TAVILY_KEY     = os.getenv("TAVILY_API_KEY")
+if os.getenv("LANGCHAIN_API_KEY"):
+    os.environ.setdefault("LANGCHAIN_TRACING_V2", "true")
+    from langchain_community.utils import configure_langsmith
+    configure_langsmith(project_name="gaia-agent")
 
 # ---------------------------------------------------------------------
-# Fehler-Wrapper – behält Doc-String dank wraps
+# 1) Helfer:  Fehler-Decorator + Backoff-Wrapper
 # ---------------------------------------------------------------------
-import functools
 def error_guard(fn):
+    """Fängt Tool-Fehler ab & gibt String zurück (bricht Agent nicht ab)."""
     @functools.wraps(fn)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args, **kw):
         try:
-            return fn(*args, **kwargs)
+            return fn(*args, **kw)
         except Exception as e:
             return f"ERROR: {e}"
     return wrapper
 
 
-# ---------------------------------------------------------------------
-# 1) fetch_gaia_file   – Datei vom GAIA-Server holen
-# ---------------------------------------------------------------------
-GAIA_FILE_ENDPOINT = "https://agents-course-unit4-scoring.hf.space/file"
-
-@tool
-@error_guard
-def fetch_gaia_file(task_id: str) -> str:
-    """Download the attachment for the given GAIA task_id and return local path."""
-    url = f"{GAIA_FILE_ENDPOINT}/{task_id}"
-    try:
-        response = requests.get(url, timeout=30)
-        response.raise_for_status()
-        file_name = response.headers.get("x-gaia-filename", f"{task_id}")
-        tmp_path = tempfile.gettempdir() + "/" + file_name
-        with open(tmp_path, "wb") as f:
-            f.write(response.content)
-        return tmp_path
-    except Exception as e:
-        return f"ERROR: could not fetch file – {e}"
+def with_backoff(fn, tries: int = 4, delay: int = 4):
+    """Synchrones Retry-Wrapper für LLM-Aufrufe."""
+    for t in range(tries):
+        try:
+            return fn()
+        except Exception as e:
+            if ("429" in str(e) or "RateLimit" in str(e)) and t < tries - 1:
+                time.sleep(delay)
+                delay *= 2
+                continue
+            raise
 
 # ---------------------------------------------------------------------
-# 2) CSV-Parser
+# 2) Eigene Tools  (CSV / Excel)
 # ---------------------------------------------------------------------
-import pandas as pd
-
 @tool
 @error_guard
 def parse_csv(file_path: str, query: str = "") -> str:
-    """Load a CSV file and answer a quick pandas query (optional)."""
+    """Load a CSV file and (optional) run a pandas query."""
     df = pd.read_csv(file_path)
     if not query:
-        return f"Loaded CSV with {len(df)} rows and {len(df.columns)} cols.\nColumns: {list(df.columns)}"
+        return f"Rows={len(df)}, Cols={list(df.columns)}"
     try:
-        result = df.query(query)
-        return result.to_markdown()
+        return df.query(query).to_markdown(index=False)
     except Exception as e:
-        return f"ERROR in pandas query: {e}"
+        return f"ERROR query: {e}"
+
 
-# ---------------------------------------------------------------------
-# 3) Excel-Parser
-# ---------------------------------------------------------------------
 @tool
 @error_guard
-def parse_excel(file_path: str, query: str = "") -> str:
-    """Load an Excel file (first sheet) and answer a pandas query (optional)."""
-    df = pd.read_excel(file_path)
+def parse_excel(file_path: str, sheet: str | int | None = None, query: str = "") -> str:
+    """Load an Excel sheet (name or index) and (optional) run a pandas query."""
+    sheet_arg = int(sheet) if isinstance(sheet, str) and sheet.isdigit() else sheet or 0
+    df = pd.read_excel(file_path, sheet_name=sheet_arg)
     if not query:
-        return f"Loaded Excel with {len(df)} rows and {len(df.columns)} cols.\nColumns: {list(df.columns)}"
+        return f"Rows={len(df)}, Cols={list(df.columns)}"
     try:
-        result = df.query(query)
-        return result.to_markdown()
+        return df.query(query).to_markdown(index=False)
     except Exception as e:
-        return f"ERROR in pandas query: {e}"
+        return f"ERROR query: {e}"
 
 # ---------------------------------------------------------------------
-# 4) Gemini-Audio-Transkription
+# 3) Externe Search-Tools (Tavily, Wikipedia)
 # ---------------------------------------------------------------------
 @tool
 @error_guard
-def gemini_transcribe_audio(file_path: str, prompt: str = "Transcribe the audio.") -> str:
-    """Use Gemini to transcribe an audio file."""
-    with open(file_path, "rb") as f:
-        b64 = base64.b64encode(f.read()).decode()
-    mime = mimetypes.guess_type(file_path)[0] or "audio/mpeg"
-    message = HumanMessage(
-        content=[
-            {"type": "text", "text": prompt},
-            {"type": "media", "data": b64, "mime_type": mime},
-        ]
-    )
-    resp = asyncio.run(safe_invoke([message]))
-    return resp.content if hasattr(resp, "content") else str(resp)
-
-# ---------------------------------------------------------------------
-# 5) Bild-Beschreibung
-# ---------------------------------------------------------------------
-@tool
-@error_guard
-def describe_image(file_path: str, prompt: str = "Describe this image.") -> str:
-    """Gemini vision – Bild beschreiben."""
-    from PIL import Image
-    img = Image.open(file_path)
-    message = HumanMessage(
-        content=[
-            {"type": "text", "text": prompt},
-            img,  # langchain übernimmt Encoding
-        ]
-    )
-    resp = asyncio.run(safe_invoke([message]))
-    return resp.content
+def web_search(query: str, max_results: int = 5) -> str:
+    """Search the web via Tavily and return markdown list of results."""
+    api_key = os.getenv("TAVILY_API_KEY")
+    hits = TavilySearchResults(max_results=max_results, api_key=api_key).invoke(query)
+    if not hits:
+        return "No results."
+    return "\n".join(f"{h['title']} – {h['url']}" for h in hits)
 
-# ---------------------------------------------------------------------
-# 6) OCR-Tool
-# ---------------------------------------------------------------------
-@tool
-@error_guard
-def ocr_image(file_path: str, lang: str = "eng") -> str:
-    """Extract text from an image via pytesseract."""
-    try:
-        import pytesseract
-        from PIL import Image
-        text = pytesseract.image_to_string(Image.open(file_path), lang=lang)
-        return text.strip() or "No text found."
-    except Exception as e:
-        return f"ERROR: {e}"
 
-# ---------------------------------------------------------------------
-# 7) Tavily-Web-Suche
-# ---------------------------------------------------------------------
 @tool
 @error_guard
-def web_search(query: str, max_results: int = 5) -> str:
-    """Search the web via Tavily and return a markdown list of results."""
-    hits = TavilySearchResults(max_results=max_results, api_key=TAVILY_KEY).invoke(query)
-    if not hits:
-        return "No results."
-    return "\n\n".join(f"{h['title']} – {h['url']}" for h in hits)
+def wiki_search(query: str, sentences: int = 3) -> str:
+    """Quick Wikipedia summary."""
+    wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=4000)
+    res = wrapper.run(query)
+    return "\n".join(res.split(". ")[:sentences]) if res else "No article found."
 
 # ---------------------------------------------------------------------
-# 8) Kleiner Rechner
+# 4) Python-REPL Tool (fertig aus LangChain)
 # ---------------------------------------------------------------------
-@tool
-@error_guard
-def simple_calculator(operation: str, a: float, b: float) -> float:
-    """Basic maths (add, subtract, multiply, divide)."""
-    ops = {
-        "add": a + b,
-        "subtract": a - b,
-        "multiply": a * b,
-        "divide": a / b if b else float("inf"),
-    }
-    return ops.get(operation, f"ERROR: unknown op '{operation}'")
+python_repl = PythonAstREPLTool()
 
 # ---------------------------------------------------------------------
-# LLM + Semaphore-Throttle  (Gemini 2.0 Flash)
+# 5) LLM – Gemini Flash, an Tools gebunden
 # ---------------------------------------------------------------------
 gemini_llm = ChatGoogleGenerativeAI(
+    google_api_key=os.getenv("GOOGLE_API_KEY"),
     model="gemini-2.0-flash",
-    google_api_key=GOOGLE_API_KEY,
     temperature=0,
     max_output_tokens=2048,
-).bind_tools([
-    fetch_gaia_file, parse_csv, parse_excel,
-    gemini_transcribe_audio, describe_image, ocr_image,
-    web_search, simple_calculator,] ,return_named_tools=True)
-
-LLM_SEMA = asyncio.Semaphore(2)           # 3 gleichz. Anfragen ≈ < 15/min
-
-# safe_invoke neu (ersetzt die alte Funktion)
-async def safe_invoke(msgs, tries: int = 4):
-    """Gemini-Aufruf mit Semaphor + Exponential-Back-off bei 429 / Netzfehlern."""
-    delay = 4
-    for t in range(tries):
-        async with LLM_SEMA:
-            try:
-                return await gemini_llm.ainvoke(msgs)
-            except Exception as e:
-                # nur bei Rate-Limit oder Netzwerk erneut versuchen
-                if ("429" in str(e) or "RateLimit" in str(e)) and t < tries - 1:
-                    await asyncio.sleep(delay)
-                    delay *= 2          # 4 s, 8 s, 16 s …
-                    continue
-                raise
-
-# ---------------------------------------------------------------------
-# System-Prompt
-# ---------------------------------------------------------------------
-system_prompt = SystemMessage(content="""
-You are a helpful assistant tasked with answering questions using a set of tools. 
-Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
-FINAL ANSWER: [YOUR FINAL ANSWER]. 
-YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither 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.
-Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
-""")
+).bind_tools(
+    [web_search, wiki_search, parse_csv, parse_excel, python_repl],
+    return_named_tools=True,
+)
+
+# ---------------------------------------------------------------------
+# 6) System-Prompt (ReAct, keine Prefixe im Final-Output!)
+# ---------------------------------------------------------------------
+SYSTEM_PROMPT = SystemMessage(
+    content=(
+        "You are a helpful assistant with access to Python tools.\n"
+        "• Think step by step.\n"
+        "• Call a tool when needed – reply in this JSON format:\n"
+        "  {\"tool\": \"<tool_name>\", \"tool_input\": { ... }}\n"
+        "• When you have the answer, reply with the answer **only** "
+        "– no prefix, no explanations.\n"
+        "Answer format rules:\n"
+        "  • Single number → no separators / units unless required.\n"
+        "  • Single string → no articles/abbrev.\n"
+        "  • List  → comma + single space separated, keep required order.\n"
+    )
+)
 
 # ---------------------------------------------------------------------
-# LangGraph – Assistant-Node
+# 7) LangGraph – Planner + Tools + Router
 # ---------------------------------------------------------------------
-def assistant(state: MessagesState):
+def planner(state: MessagesState):
+    """LLM-Planner – entscheidet, ob Tool nötig oder Final Answer erreicht."""
     msgs = state["messages"]
     if msgs[0].type != "system":
-        msgs = [system_prompt] + msgs
-    resp = asyncio.run(safe_invoke(msgs))
-    finished = resp.content.lower().lstrip().startswith("final answer") or not resp.tool_calls
+        msgs = [SYSTEM_PROMPT] + msgs
+    resp = with_backoff(lambda: gemini_llm.invoke(msgs))
+    finished = (
+        not getattr(resp, "tool_calls", None)  # keine Toolaufrufe
+        and "\n" not in resp.content          # heuristik: kurze Endantwort
+    )
     return {"messages": [resp], "should_end": finished}
 
 def route(state):
     return "END" if state["should_end"] else "tools"
 
+# Tool-Knoten
+TOOLS = [web_search, wiki_search, parse_csv, parse_excel, python_repl]
+
+graph = StateGraph(MessagesState)
+graph.add_node("planner", planner)
+graph.add_node("tools", ToolNode(TOOLS))
+graph.add_edge(START, "planner")
+graph.add_conditional_edges("planner", route, {"tools": "tools", "END": END})
+
+# compile → LangGraph-Executor
+agent_executor = graph.compile(max_iterations=8)
+
 # ---------------------------------------------------------------------
-# Tools-Liste & Graph
+# 8) Öffentliche Klasse  –  wird von app.py / logic.py verwendet
 # ---------------------------------------------------------------------
-tools = [
-    fetch_gaia_file, parse_csv, parse_excel,
-    gemini_transcribe_audio, describe_image, ocr_image,
-    web_search, simple_calculator,
-]
+class GaiaAgent:
+    """LangChain·LangGraph-Agent für GAIA Level 1."""
 
-builder = StateGraph(MessagesState)
-builder.add_node("assistant", assistant)
-builder.add_node("tools", ToolNode(tools))
-builder.add_edge(START, "assistant")
-builder.add_conditional_edges("assistant", route, {"tools": "tools", "END": END})
+    def __init__(self):
+        print("✅ GaiaAgent initialised (LangGraph)")
 
-# Compile
-agent_executor = builder.compile()
+    def __call__(self, task_id: str, question: str) -> str:
+        """Run the agent on a single GAIA question → exact answer string."""
+        start_state = {"messages": [HumanMessage(content=question)]}
+        final_state = agent_executor.invoke(start_state)
+        # letze Message enthält Antwort
+        answer = final_state["messages"][-1].content
+        return answer.strip()