Update agent.py

agent.py

@@ -1,23 +1,29 @@
 import os
 from pathlib import Path
 from typing import Optional, Union
+
 import pandas as pd
 from dotenv import load_dotenv
+from tabulate import tabulate
+
 from smolagents import (
-    CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool,
-    LiteLLMModel, PythonInterpreterTool, WikipediaSearchTool
+    CodeAgent,
+    DuckDuckGoSearchTool,
+    FinalAnswerTool,
+    LiteLLMModel,
+    PythonInterpreterTool,
+    WikipediaSearchTool,
 )
 from smolagents.tools import Tool
-from tabulate import tabulate
 
-# Load environment variables (make sure GEMINI_API_KEY is set in your Space secrets)
+# Load environment variables
 load_dotenv()
 
+# ✅ Initialize the model with Google AI Studio Gemini
 model = LiteLLMModel(
-    model_id="gemini-2.5-flash-preview-09-2025",
-    api_key=os.getenv("GEMINI_API_KEY"),
-    provider="google_ai",
-
+    model_id="gemini-2.5-flash-preview-09-2025",   # short name only
+    api_key=os.getenv("GEMINI_API_KEY"),          # set in Hugging Face secrets
+    provider="google_ai",                         # force AI Studio, not Vertex
 )
 
 
@@ -55,9 +61,7 @@ class ExcelToTextTool(Tool):
 
         try:
             sheet: Union[str, int] = (
-                int(sheet_name)
-                if sheet_name and sheet_name.isdigit()
-                else sheet_name or 0
+                int(sheet_name) if sheet_name and sheet_name.isdigit() else sheet_name or 0
             )
 
             df = pd.read_excel(file_path, sheet_name=sheet)
@@ -75,7 +79,7 @@ class GaiaAgent:
     """An agent capable of using tools to answer general questions."""
 
     def __init__(self):
-        print("GaiaAgent initialized with tools.")
+        print("✅ GaiaAgent initialized with tools.")
 
         tools = [
             DuckDuckGoSearchTool(),
@@ -93,12 +97,12 @@ class GaiaAgent:
         )
 
     def __call__(self, task_id: str, question: str) -> str:
-        print(f"Agent received task_id='{task_id}' | question='{question[:50]}...'")
+        print(f"🔹 Agent received task_id='{task_id}' | question='{question[:50]}...'")
         try:
             answer = self.agent.run(question)
             if not answer:
                 answer = "⚠️ Sorry, I could not generate a valid response."
         except Exception as e:
             answer = f"⚠️ Agent failed with error: {e}"
-        print(f"Agent returning answer: {answer}")
+        print(f"🔹 Agent returning answer: {answer}")
         return answer