Update agent.py

agent.py

@@ -49,7 +49,7 @@ def _download_file(file_id: str) -> bytes:
 # --------------------------------------------------------------------------- #
 class GeminiModel:
     """
-    Thin adapter around google-genai.Client so it can be used by smolagents.
+    Thin adapter around google-genai Client for smolagents.
     """
 
     def __init__(
@@ -61,15 +61,15 @@ class GeminiModel:
         api_key = os.getenv("GOOGLE_API_KEY")
         if not api_key:
             raise EnvironmentError("GOOGLE_API_KEY is not set.")
-        # One client per process is enough
         self.client = genai.Client(api_key=api_key)
         self.model_name = model_name
         self.temperature = temperature
         self.max_tokens = max_tokens
 
-    # ---------- Text-only convenience ---------- #
+    # ---------- main generation helpers ---------- #
     def call(self, prompt: str, **kwargs) -> str:
-        response = self.client.models.generate_content(
+        """Text-only helper used by __call__."""
+        resp = self.client.models.generate_content(
             model=self.model_name,
             contents=prompt,
             generation_config=gtypes.GenerateContentConfig(
@@ -77,23 +77,16 @@ class GeminiModel:
                 max_output_tokens=self.max_tokens,
             ),
         )
-        return response.text.strip()
+        return resp.text.strip()
 
-    # ---------- smolagents will use this when messages are present ---------- #
     def call_messages(self, messages, **kwargs) -> str:
-        """
-        `messages` is a list of dictionaries with keys 'role' | 'content'.
-        If `content` is already a list[types.Content], we forward it as-is.
-        Otherwise we concatenate to a single string prompt.
-        """
-        sys_msg, user_msg = messages  # CodeAgent always sends two
-        if isinstance(user_msg["content"], list):
-            # Multimodal path – pass system text first, then structured user parts
-            contents = [sys_msg["content"], *user_msg["content"]]
-        else:
-            # Text prompt path
-            contents = f"{sys_msg['content']}\n\n{user_msg['content']}"
-        response = self.client.models.generate_content(
+        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']}"
+        )
+        resp = self.client.models.generate_content(
             model=self.model_name,
             contents=contents,
             generation_config=gtypes.GenerateContentConfig(
@@ -101,7 +94,11 @@ class GeminiModel:
                 max_output_tokens=self.max_tokens,
             ),
         )
-        return response.text.strip()
+        return resp.text.strip()
+
+    # ---------- make the instance itself callable ---------- #
+    def __call__(self, prompt: str, **kwargs) -> str:  # <-- NEW
+        return self.call(prompt, **kwargs)
 
 
 # --------------------------------------------------------------------------- #