Update app.py

app.py

@@ -1,11 +1,11 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent,DuckDuckGoSearchTool, LiteLLMModel,load_tool,tool
 import datetime
 import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
-
-# from Gradio_UI import GradioUI
+import os
+from Gradio_UI import GradioUI
 
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
 @tool
@@ -39,12 +39,30 @@ final_answer = FinalAnswerTool()
 # If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
 # model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 
 
-model = HfApiModel(
-max_tokens=2096,
-temperature=0.5,
-model_id='google/gemma-2b-it',# it is possible that this model may be overloaded
-custom_role_conversions=None,
-)
+os.environ["GOOGLE_API_KEY"] = "AIzaSyBcJrlnDDdWtjUDiLrisSOPuaAGizCLKO4"
+gemini_api_key = os.environ.get("GOOGLE_API_KEY")
+
+try:
+    # LiteLLM uses 'gemini/' prefix for Google AI Studio models
+    gemini_model = LiteLLMModel(
+        model_id="gemini/gemini-1.5-flash-latest",
+        api_key=gemini_api_key,
+        temperature = 0.5,
+        max_tokens = 2096,
+        custom_role_conversions=None
+    )
+    print("Successfully initialized LiteLLMModel for Gemini 1.5 Flash.")
+
+except Exception as e:
+    print(f"Failed to initialize LiteLLMModel: {e}")
+    gemini_model = None
+    
+# model = HfApiModel(
+# max_tokens=2096,
+# temperature=0.5,
+# model_id='google/gemma-2b-it',# it is possible that this model may be overloaded
+# custom_role_conversions=None,
+# )
 
 
 # Import tool from Hub
@@ -54,7 +72,7 @@ with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
     
 agent = CodeAgent(
-    model=model,
+    model=gemini_model,
     tools=[final_answer], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
@@ -66,11 +84,4 @@ agent = CodeAgent(
 )
 
 
-# GradioUI(agent).launch()
-try:
-    result = agent.run("What weather is it in Europe/London?") # Or another test message
-    print("Agent Result:", result)
-except Exception as e:
-    print("An error occurred during agent run:", e)
-    import traceback
-    traceback.print_exc() # <--- This will print the full error
+GradioUI(agent).launch()