Switch agent model to Gemini via LiteLLM

app.py

@@ -11,6 +11,7 @@ import pandas as pd
 from smolagents import (
     CodeAgent,
     DuckDuckGoSearchTool,
+    LiteLLMModel,
     InferenceClientModel,
     VisitWebpageTool,
     tool,
@@ -210,40 +211,41 @@ class BasicAgent:
     """Same name and call interface as the template: agent(question) -> str."""
 
     def __init__(self):
-        token = os.getenv("HF_TOKEN")
-        if not token:
-            raise RuntimeError("HF_TOKEN is not set. Add it as a Space *secret* (or export it locally).")
-
-        self.model = GuardedInferenceClientModel(
-            model_id=MODEL_ID,
-            token=token,
-            provider=HF_PROVIDER,
+        # Fetch Google Gemini API Key from Space secrets
+        gemini_key = os.getenv("GEMINI_API_KEY")
+        if not gemini_key:
+            raise RuntimeError("GEMINI_API_KEY is not set. Add it as a Space *secret*.")
+
+        # Initialize Gemini 2.5 Flash / 1.5 Flash using LiteLLMModel wrapper
+        self.model = LiteLLMModel(
+            model_id="gemini/gemini-2.5-flash",  # or "gemini/gemini-1.5-flash"
+            api_key=gemini_key,
             temperature=0.1,
-            max_tokens=1500,
         )
 
         visit = VisitWebpageTool()
-        visit.max_output_length = MAX_PAGE_CHARS   # honored by recent versions; harmless otherwise
+        visit.max_output_length = MAX_PAGE_CHARS
 
         base = dict(
             tools=[RobustDuckDuckGoSearchTool(max_results=6, rate_limit=2.0), visit, read_file],
             model=self.model,
-            additional_authorized_imports=["pandas", "numpy", "json", "csv", "io",
-                                           "string", "fractions", "decimal", "openpyxl"],
+            additional_authorized_imports=[
+                "pandas", "numpy", "json", "csv", "io",
+                "string", "fractions", "decimal", "openpyxl"
+            ],
             max_steps=MAX_STEPS,
             verbosity_level=1,
         )
-        # Answer rules go in the system prompt when this smolagents version supports
-        # `instructions`; otherwise they are prepended to every task.
+
         try:
             self.agent = CodeAgent(**base, instructions=ANSWER_RULES, max_print_outputs_length=4000)
             self.rules_in_system = True
         except TypeError:
             self.agent = CodeAgent(**base)
             self.rules_in_system = False
-        print(f"BasicAgent initialized | model={MODEL_ID} | provider={HF_PROVIDER or 'auto'} | "
-              f"max_steps={MAX_STEPS} | rules_in_system={self.rules_in_system}")
-
+            
+        print(f"BasicAgent initialized with Gemini Flash | max_steps={MAX_STEPS}")
+        
     def _build_task(self, question: str, file_name: str, file_path: str | None) -> str:
         parts = [question.strip()]
         if file_name and file_path:

requirements.txt

@@ -9,6 +9,8 @@ huggingface_hub
 spaces
 markdownify
 smolagents[toolkit]
+litellm
+google-generativeai
 # Which DuckDuckGo package smolagents imports depends on its version.
 # Installing both is harmless.
 ddgs