Update agent.py

agent.py

@@ -15,14 +15,30 @@ from tabulate import tabulate
 # Load environment variables
 load_dotenv()
 
-# Initialize the model with OpenRouter Gemini
+# ----------------------------
+# 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
 model = LiteLLMModel(
-    model_id=f"openrouter/{os.getenv('GEMINI_MODEL', 'google/gemini-2.5-flash-preview-09-2025')}",
-    api_key=os.getenv("GEMINI_API_KEY"),   # your OpenRouter API key
-    base_url="https://openrouter.ai/api/v1"  # required for OpenRouter
+    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}
 )
 
-
+# ----------------------------
+# Custom tool: Excel to Markdown
+# ----------------------------
 class ExcelToTextTool(Tool):
     """Render an Excel worksheet as a Markdown table."""
 
@@ -72,6 +88,9 @@ 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."""
 
@@ -98,11 +117,17 @@ class GaiaAgent:
         try:
             answer = self.agent.run(question)
         except Exception as e:
-            answer = f"⚠️ Agent failed: {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 could not generate a valid response."
+            answer = "⚠️ Sorry, I couldn't generate a response."
 
         print(f"Agent returning answer: {answer}")
         return answer
-