Update agent.py

agent.py

@@ -1,8 +1,6 @@
 import os
 from pathlib import Path
 from typing import Optional, Union
-import smolagents
-
 import pandas as pd
 from dotenv import load_dotenv
 from smolagents import (
@@ -12,33 +10,16 @@ from smolagents import (
 from smolagents.tools import Tool
 from tabulate import tabulate
 
-# Load environment variables
+# Load environment variables (make sure GEMINI_API_KEY is set in your Space secrets)
 load_dotenv()
 
-# ----------------------------
-# Initialize the Google Gemini LLM
-# ----------------------------
-model_name = os.getenv("GEMINI_MODEL", "google/gemini-2.5-flash-preview-09-2025")
-g_api_key = os.getenv("GEMINI_API_KEY")
-vertex_project = os.getenv("VERTEX_PROJECT")
-vertex_location = os.getenv("VERTEX_LOCATION", "us-central1")
-
-if not g_api_key:
-    raise RuntimeError(
-        "Google API key not found. Set GEMINI_API_KEY in your Space secrets."
-    )
-
-# LiteLLM expects provider prefix
+# Initialize the Gemini model using Google AI Studio API key
 model = LiteLLMModel(
-    model_id=f"vertex_ai/gemini/{model_name}",
-    api_key=g_api_key,
-    # Optional: pass project/location if supported by smolagents
-    # provider_config={"vertex_project": vertex_project, "vertex_location": vertex_location}
+    model_id="vertex_ai/gemini/google/gemini-2.5-flash-preview-09-2025",
+    api_key=os.getenv("GEMINI_API_KEY"),
 )
 
-# ----------------------------
-# Custom tool: Excel to Markdown
-# ----------------------------
+
 class ExcelToTextTool(Tool):
     """Render an Excel worksheet as a Markdown table."""
 
@@ -66,6 +47,7 @@ class ExcelToTextTool(Tool):
     output_type = "string"
 
     def forward(self, excel_path: str, sheet_name: Optional[str] = None) -> str:
+        """Load the Excel file and return the sheet as a Markdown table."""
         file_path = Path(excel_path).expanduser().resolve()
         if not file_path.is_file():
             return f"Error: Excel file not found at {file_path}"
@@ -88,9 +70,6 @@ class ExcelToTextTool(Tool):
             return f"Error reading Excel file: {e}"
 
 
-# ----------------------------
-# GaiaAgent: handles multiple tools
-# ----------------------------
 class GaiaAgent:
     """An agent capable of using tools to answer general questions."""
 
@@ -116,18 +95,9 @@ class GaiaAgent:
         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:
-            err_text = str(e)
-            # Catch common LLM errors
-            if "Insufficient credits" in err_text:
-                answer = "⚠️ LLM request failed: insufficient Google AI Studio credits."
-            elif "Provider NOT provided" in err_text or "vertex" in err_text.lower():
-                answer = f"⚠️ LLM provider error: {err_text}"
-            else:
-                answer = f"⚠️ Agent failed: {err_text}"
-
-        if not answer:
-            answer = "⚠️ Sorry, I couldn't generate a response."
-
+            answer = f"⚠️ Agent failed with error: {e}"
         print(f"Agent returning answer: {answer}")
         return answer