Update agent.py

agent.py

@@ -1,24 +1,22 @@
 """
-agent.py ­– Gemini-smolagents baseline using google-genai SDK
-------------------------------------------------------------
-Environment variables
----------------------
-GOOGLE_API_KEY            API key from Google AI Studio
-Optional:
-    GAIA_API_URL          GAIA evaluation endpoint (default: official URL)
-
-This file defines:
-    • GeminiModel  – wraps google-genai for smolagents
-    • gaia_file_reader  – custom tool to fetch <file:xyz> attachments
-    • GeminiAgent  – CodeAgent with Python / Search / File tools + Gemini model
+agent.py – Gemini-smolagents baseline using google-genai SDK
+-----------------------------------------------------------
+Environment
+-----------
+GOOGLE_API_KEY   – API key from Google AI Studio
+GAIA_API_URL     – (optional) override for the GAIA scoring endpoint
 """
 
-import os
-import re
+from __future__ import annotations
+
 import base64
 import mimetypes
-import requests
+import os
+import re
+from typing import List
+
 import google.genai as genai
+import requests
 from google.genai import types as gtypes
 from smolagents import (
     CodeAgent,
@@ -28,13 +26,12 @@ from smolagents import (
 )
 
 # --------------------------------------------------------------------------- #
-# Constants & helpers
+# constants & helpers
 # --------------------------------------------------------------------------- #
 DEFAULT_API_URL = os.getenv(
     "GAIA_API_URL", "https://agents-course-unit4-scoring.hf.space"
 )
-FILE_TAG = re.compile(r"<file:([^>]+)>")
-
+FILE_TAG = re.compile(r"<file:([^>]+)>")  # <file:xyz>
 
 def _download_file(file_id: str) -> bytes:
     """Download the attachment for a GAIA task."""
@@ -43,13 +40,13 @@ def _download_file(file_id: str) -> bytes:
     resp.raise_for_status()
     return resp.content
 
-
 # --------------------------------------------------------------------------- #
-# Model wrapper
+# model wrapper
 # --------------------------------------------------------------------------- #
 class GeminiModel:
     """
-    Thin adapter around google-genai Client for smolagents.
+    Minimal adapter around google-genai.Client so the instance itself is
+    callable (required by smolagents).
     """
 
     def __init__(
@@ -62,47 +59,44 @@ class GeminiModel:
         if not api_key:
             raise EnvironmentError("GOOGLE_API_KEY is not set.")
         self.client = genai.Client(api_key=api_key)
+
         self.model_name = model_name
         self.temperature = temperature
         self.max_tokens = max_tokens
 
-    # ---------- main generation helpers ---------- #
-    def call(self, prompt: str, **kwargs) -> str:
-        """Text-only helper used by __call__."""
-        resp = self.client.models.generate_content(
-            model=self.model_name,
-            contents=prompt,
-            config=gtypes.GenerateContentConfig(
-                temperature=self.temperature,
-                max_output_tokens=self.max_tokens,
-            ),
-        )
-        return resp.text.strip()
-
-    def call_messages(self, messages, **kwargs) -> str:
-        sys_msg, user_msg = messages
-        contents = (
-            [sys_msg["content"], *user_msg["content"]]
-            if isinstance(user_msg["content"], list)
-            else f"{sys_msg['content']}\n\n{user_msg['content']}"
-        )
+    # internal helper -------------------------------------------------------- #
+    def _genai_call(
+        self,
+        contents,
+        system_instruction: str,
+    ) -> str:
         resp = self.client.models.generate_content(
             model=self.model_name,
             contents=contents,
             config=gtypes.GenerateContentConfig(
+                system_instruction=system_instruction,
                 temperature=self.temperature,
                 max_output_tokens=self.max_tokens,
             ),
         )
         return resp.text.strip()
 
-    # ---------- make the instance itself callable ---------- #
-    def __call__(self, prompt: str, **kwargs) -> str:  # <-- NEW
-        return self.call(prompt, **kwargs)
+    # public helpers --------------------------------------------------------- #
+    def __call__(self, prompt: str, system_instruction: str, **__) -> str:
+        """Used by CodeAgent for plain-text questions."""
+        return self._genai_call(prompt, system_instruction)
 
+    def call_parts(
+        self,
+        parts: List[gtypes.Part],
+        system_instruction: str,
+    ) -> str:
+        """Multimodal path used by GeminiAgent for <file:…> questions."""
+        user_content = gtypes.Content(role="user", parts=parts)
+        return self._genai_call([user_content], system_instruction)
 
 # --------------------------------------------------------------------------- #
-# Custom tool: fetch GAIA attachments
+# custom tool: fetch GAIA attachments
 # --------------------------------------------------------------------------- #
 @tool
 def gaia_file_reader(file_id: str) -> str:
@@ -110,12 +104,11 @@ def gaia_file_reader(file_id: str) -> str:
     Download a GAIA attachment and return its contents.
 
     Args:
-        file_id: The identifier that appears inside a <file:...> placeholder
-                 in the GAIA question prompt.
+        file_id: identifier that appears inside a <file:...> placeholder.
 
     Returns:
-        A base-64 string for binary files (images, PDF, etc.) or UTF-8 text for
-        plain-text files.
+        base64-encoded string for binary files (images, PDFs, …) or decoded
+        UTF-8 text for textual files.
     """
     try:
         raw = _download_file(file_id)
@@ -123,47 +116,49 @@ def gaia_file_reader(file_id: str) -> str:
         if mime.startswith("text") or mime in ("application/json",):
             return raw.decode(errors="ignore")
         return base64.b64encode(raw).decode()
-    except Exception as exc:
+    except Exception as exc:  # pragma: no cover
         return f"ERROR downloading {file_id}: {exc}"
 
-
 # --------------------------------------------------------------------------- #
-# Final agent class
+# final agent
 # --------------------------------------------------------------------------- #
 class GeminiAgent:
-    def __init__(self):
-        self.system_prompt = (
-            "You are a concise, highly accurate assistant. "
-            "Unless explicitly required, reply with ONE short sentence. "
-            "Use the provided tools if needed. "
-            "All answers are graded by exact string match."
-        )
+    """Instantiated once in app.py; called once per question."""
 
-        model = GeminiModel()
-        tools = [
-            PythonInterpreterTool(),
-            DuckDuckGoSearchTool(),
-            gaia_file_reader,
-        ]
+    SYSTEM_PROMPT = (
+        "You are a concise, highly accurate assistant. "
+        "Unless explicitly required, reply with ONE short sentence. "
+        "Use the provided tools if needed. "
+        "All answers are graded by exact string match."
+    )
 
-        # ✨  system_prompt removed – newest smolagents doesn't take it
+    def __init__(self):
+        self.model = GeminiModel()
         self.agent = CodeAgent(
-            model=model,
-            tools=tools,
-            # any other kwargs (executor_type, additional_authorized_imports…)
+            model=self.model,
+            tools=[
+                PythonInterpreterTool(),
+                DuckDuckGoSearchTool(),
+                gaia_file_reader,
+            ],
             verbosity_level=0,
         )
-        print("✅ GeminiAgent ready.")
+        print("✅ GeminiAgent initialised.")
 
+    # --------------------------------------------------------------------- #
+    # main entry point
+    # --------------------------------------------------------------------- #
     def __call__(self, question: str) -> str:
         file_ids = FILE_TAG.findall(question)
 
-        # -------- multimodal branch -------- #
+        # ---------- multimodal branch (images / files) -------------------- #
         if file_ids:
-            parts: list[gtypes.Part] = []
+            parts: List[gtypes.Part] = []
+
             text_part = FILE_TAG.sub("", question).strip()
             if text_part:
                 parts.append(gtypes.Part.from_text(text_part))
+
             for fid in file_ids:
                 try:
                     img_bytes = _download_file(fid)
@@ -172,19 +167,14 @@ class GeminiAgent:
                         gtypes.Part.from_bytes(data=img_bytes, mime_type=mime)
                     )
                 except Exception as exc:
-                    parts.append(
-                        gtypes.Part.from_text(f"[FILE {fid} ERROR: {exc}]")
-                    )
-            messages = [
-                {"role": "system", "content": self.system_prompt},
-                {"role": "user", "content": parts},
-            ]
-            answer = self.agent.model.call_messages(messages)
+                    parts.append(gtypes.Part.from_text(f"[FILE {fid} ERROR: {exc}]"))
+
+            answer = self.model.call_parts(parts, system_instruction=self.SYSTEM_PROMPT)
 
-        # -------- text-only branch -------- #
+        # ---------- plain-text branch ------------------------------------- #
         else:
-            # prepend system prompt to the user question
-            full_prompt = f"{self.system_prompt}\n\n{question}"
-            answer = self.agent(full_prompt)
+            # Prepend system prompt to make sure CodeAgent->model sees it.
+            prompt = f"{self.SYSTEM_PROMPT}\n\n{question}"
+            answer = self.agent.model(prompt, system_instruction=self.SYSTEM_PROMPT)
 
         return answer.rstrip(" .\n\r\t")