Update app.py

app.py

@@ -1,10 +1,11 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel, load_tool, tool, LiteLLMModel
 import datetime
 import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
 from typing import Tuple
+import os
 
 from Gradio_UI import GradioUI
 
@@ -47,13 +48,19 @@ def get_current_time_in_timezone(timezone: str) -> str:
 
 
 final_answer = FinalAnswerTool()
-model = HfApiModel(
-max_tokens=2096,
-temperature=0.5,
-model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
-custom_role_conversions=None,
-)
+# model = HfApiModel(
+# max_tokens=2096,
+# temperature=0.5,
+# model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
+# custom_role_conversions=None,
+# )
 
+model = LiteLLMModel(
+    model_id = "gemini/gemini-2.0-flash", 
+    api_key = os.getenv("GEMINI_API_KEY"),
+    temperature = 1,
+    max_tokens=4096
+)
 
 # Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)