Update agent.py

agent.py

@@ -1,15 +1,17 @@
-from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, HfApiModel # OpenAIServerModel
+from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, LiteLLMModel # HfApiModel, OpenAIServerModel
 import asyncio
 
 class MagAgent:
     def __init__(self):
         """Initialize the MagAgent with search tools."""
         print("Initializing MagAgent with search tools...")
-        model = HfApiModel(
+        model = LiteLLMModel(
+            model_id="gemini/gemini-2.0-flash",
+            api_key=os.getenv("GEMINI_KEY"),
+            max_tokens=8192
+#        model = HfApiModel(
 #            max_tokens=2096,
 #            temperature=0.5,
-#        #    model_id="gemini/gemini-2.0-flash-lite-001",
-#            # api_key = os.environ["GEMINI_KEY"], #userdata.get("GEMINI_KEY")
 #            model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
 #            custom_role_conversions=None,
         )