Update agents.py

agents.py

@@ -1,25 +1,18 @@
-from smolagents import (CodeAgent, 
-                        GradioUI, 
-                        LiteLLMModel, 
-                        OpenAIServerModel, 
-                        ChatMessage, 
-                        ToolCallingAgent)
-from smolagents.default_tools import (DuckDuckGoSearchTool, 
-                                      VisitWebpageTool, 
-                                      WikipediaSearchTool, 
+from smolagents import CodeAgent, LiteLLMModel
+from smolagents.default_tools import (DuckDuckGoSearchTool,
+                                      VisitWebpageTool,
+                                      WikipediaSearchTool,
                                       SpeechToTextTool,
                                       PythonInterpreterTool)
 import yaml
-from final_answer import FinalAnswerTool, check_reasoning, ensure_formatting
-from tools import (youtube_frames_to_images, use_vision_model, 
-                         read_file, download_file_from_url, 
-                         extract_text_from_image, analyze_csv_file, 
+from final_answer import FinalAnswerTool
+from tools import (youtube_frames_to_images, use_vision_model,
+                         read_file, download_file_from_url,
+                         extract_text_from_image, analyze_csv_file,
                          analyze_excel_file, youtube_transcribe,
-                         transcribe_audio, review_youtube_video)
+                        transcribe_audio, review_youtube_video)
 import os
-import traceback
 from dotenv import load_dotenv
-import time
 
 load_dotenv()
 
@@ -27,26 +20,13 @@ load_dotenv()
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
 
-
-    
-class SlowLiteLLMModel(LiteLLMModel):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def __call__(self, messages, **kwargs) -> ChatMessage:
-        time.sleep(15)
-        # prepend onto whatever messages the Agent built
-        return super().__call__(messages, **kwargs)
-
-
-react_model_name = "gemini-1.5-mini"
+# WORKING MODEL SETUP
 react_model = LiteLLMModel(
-    model_id="google/gemini-1.5-mini",  
+    model_id="gemini/gemini-1.5-mini",
     api_key=os.getenv("GEMINI_KEY"),
     temperature=0.2
 )
 
-
 manager_agent = CodeAgent(
     model=react_model,
     tools=[FinalAnswerTool(),
@@ -70,11 +50,8 @@ manager_agent = CodeAgent(
     planning_interval=6,
     name="Manager",
     description="The manager of the team, responsible for overseeing and guiding the team's work.",
-    # final_answer_checks=[check_reasoning, ensure_formatting],  # COMMENT THIS OUT
-    prompt_templates=prompt_templates
-)
-
-
+    prompt_templates=prompt_templates)
 
 if __name__ == "__main__":
+    from smolagents import GradioUI
     GradioUI(manager_agent).launch()