feat: gemini-model support - LiteLLM

agent.py

@@ -1,48 +1,25 @@
-from smolagents import (
-    CodeAgent,
-    DuckDuckGoSearchTool,
-    FinalAnswerTool,
-    PythonInterpreterTool,
-    VisitWebpageTool,
-    WikipediaSearchTool,
-    tool,
-)
-from smolagents.models import Model
-import google.generativeai as genai
-from dotenv import load_dotenv
 import os
 
-load_dotenv()
-
-class GeminiModel(Model):
-    def __init__(self):
-        self.api_key = os.getenv("GEMINI_API_KEY")
-        self.model_name = os.getenv("GEMINI_MODEL", "gemini-pro")
+from dotenv import load_dotenv
+from smolagents import (CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool, LiteLLMModel,
+                        PythonInterpreterTool, VisitWebpageTool, WikipediaSearchTool)
 
-        if not self.api_key:
-            raise ValueError("GEMINI_API_KEY not found in .env")
+load_dotenv()
 
-        genai.configure(api_key=self.api_key)
-        self.model = genai.GenerativeModel(self.model_name)
+model = LiteLLMModel(
+    model_id=os.getenv("GEMINI_MODEL"),
+    api_key=os.getenv("GEMINI_API_KEY")
+    )
 
-    def complete(self, prompt: str) -> str:
-        try:
-            response = self.model.generate_content(prompt)
-            return response.text.strip() if hasattr(response, "text") else str(response)
-        except Exception as e:
-            return f"Error generating content: {e}"
 
 class GaiaAgent:
     """
     An agent designed to answer questions using a combination of tools,
     including search engines, web page access, a Python interpreter, and more.
     """
-
     def __init__(self):
         print("GaiaAgent initialized with tools.")
 
-        gemini_model = GeminiModel()
-
         tools = [
             DuckDuckGoSearchTool(),
             VisitWebpageTool(),
@@ -51,10 +28,10 @@ class GaiaAgent:
             FinalAnswerTool(),
         ]
 
-        self.agent = CodeAgent(model=gemini_model, tools=tools)
+        self.agent = CodeAgent(model=model, tools=tools)
 
     def __call__(self, task_id: str, question: str) -> str:
         print(f"Agent received {task_id=}\n{question[:50]=}...")
         answer = self.agent.run(question)
         print(f"Agent returning answer: {answer}")
-        return answer
+        return answer

requirements.txt

@@ -5,4 +5,5 @@ python-dotenv
 pandas
 smolagents
 wikipedia-api
-google-generativeai
+google-generativeai
+smolagents[litellm]