Fix agents

mini_agents.py

@@ -1,8 +1,9 @@
 from smolagents import CodeAgent, InferenceClientModel
 from tools import sort_list, operate_two_numbers, convert_number, load_dataframe_from_csv
+from tools import tavily_search_tool, visit_webpage_tool
 from tools import to_dataframe, to_json, get_dataframe_data, get_dataframe_column, get_dataframe_row, get_dataframe_groupby
 from vlm_tools import download_image, image_processing, object_detection_tool, ocr_scan_tool, extract_frames_from_video
-from audio_tools import audio_to_base64, noise_reduction, audio_segmentation, speaker_diarization
+from audio_tools import transcribe_audio_tool, audio_to_base64, noise_reduction, audio_segmentation, speaker_diarization
 from community_tools import community_tools
 import os
 import logging
@@ -11,12 +12,12 @@ import yaml
 logging.basicConfig(level=logging.DEBUG)
 
 MODEL_CHOICES = {
-    "audio": ["whisper-large-v3"], 
-    "vlm": ["Qwen/Qwen2.5-VL-7B-Instruct"], 
+    "audio": ["Qwen/Qwen2.5-Coder-32B-Instruct"], 
+    "vlm": ["Qwen/Qwen2.5-Coder-32B-Instruct"], 
     "code": ["Qwen/Qwen2.5-Coder-32B-Instruct"],
     "arithmetic": ["Qwen/Qwen2.5-Coder-7B-Instruct"],
     "pandas": ["Qwen/Qwen2.5-Coder-7B-Instruct"],
-    "master": ["Qwen/Qwen2.5-Coder-32B-Instruct", "deepseek-ai/DeepSeek-V3-0324"]
+    "master": ["Qwen/Qwen2.5-Coder-32B-Instruct"]
 }
 
 with open("prompts.yaml", 'r') as stream:
@@ -37,8 +38,8 @@ audio_model = InferenceClientModel(
 
 audio_agent = CodeAgent(
     model=audio_model,
-    tools=[audio_to_base64, noise_reduction, audio_segmentation, speaker_diarization],
-    max_steps=4,
+    tools=[transcribe_audio_tool, audio_to_base64, noise_reduction, audio_segmentation, speaker_diarization],
+    max_steps=6,
     # prompt_templates=PROMPT_TEMPLATE["audio_agent"],
     additional_authorized_imports=["pydub", "pyAudioAnalysis", "base64", "io", "sklearn", "scipy", "numpy", "pandas", "json", "os", "logging", "yaml", "pyplot", "matplotlib", 'hmmlearn', 'pickle'],
     name="audio_agent",
@@ -53,7 +54,7 @@ vlm_model = InferenceClientModel(
 vlm_agent = CodeAgent(
     model=vlm_model,
     tools=[download_image, image_processing, object_detection_tool, ocr_scan_tool, extract_frames_from_video],
-    max_steps=4,
+    max_steps=6,
     # prompt_templates=PROMPT_TEMPLATE["vlm_agent"],
     additional_authorized_imports=["cv2", "numpy", "pytesseract", "requests", "base64", "onnxruntime", "PIL", "io"],
     name="vlm_agent",
@@ -68,9 +69,9 @@ arithmetic_model = InferenceClientModel(
 arithmetic_agent = CodeAgent(
     model=arithmetic_model,
     tools=[operate_two_numbers, convert_number],
-    max_steps=2,
+    max_steps=4,
     # prompt_templates=PROMPT_TEMPLATE["arithmetic_agent"],
-    additional_authorized_imports=["math", "json", "pandas", "numpy", "io", "os", "scipy", "sklearn"],
+    additional_authorized_imports=["math", "pandas", "json", "numpy", "io", "os", "logging", "yaml", "pyplot", "matplotlib", 'hmmlearn', 'pickle', 'sklearn', 'scipy'],
     name="arithmetic_agent",
     description="This agent is responsible for performing arithmetic operations on two numbers."
 )
@@ -83,9 +84,9 @@ pandas_model = InferenceClientModel(
 pandas_agent = CodeAgent(
     model=pandas_model,
     tools=[load_dataframe_from_csv, to_dataframe, to_json, get_dataframe_data, get_dataframe_column, get_dataframe_row, get_dataframe_groupby],
-    max_steps=2,
+    max_steps=4,
     # prompt_templates=PROMPT_TEMPLATE["pandas_agent"],
-    additional_authorized_imports=["math","pandas", "json", "numpy", "io", "os", "logging", "yaml", "pyplot", "matplotlib", 'hmmlearn', 'pickle'],
+    additional_authorized_imports=["math", "pandas", "json", "numpy", "io", "os", "logging", "yaml", "pyplot", "matplotlib", 'hmmlearn', 'pickle', 'sklearn', 'scipy'],
     name="pandas_agent",
     description="This agent is responsible for converting data to a dataframe, performing pandas operations on such dataframe and converting the dataframe back to a json or a csv file."
 )
@@ -123,13 +124,13 @@ master_model = InferenceClientModel(
 master_agent = CodeAgent(
     model=master_model,
     managed_agents=[audio_agent, vlm_agent, arithmetic_agent, pandas_agent],
-    tools=[sort_list, *community_tools],
+    tools=[sort_list, *community_tools, tavily_search_tool, visit_webpage_tool],
     add_base_tools=True,
-    max_steps=16,
-    additional_authorized_imports=["math","pandas", "datetime", "typing"],
-    verbosity_level=logging.DEBUG,
+    max_steps=20,
+    additional_authorized_imports=["math", "pandas", "json", "numpy", "io", "os", "logging", "yaml", "pyplot", "matplotlib", 'hmmlearn', 'pickle', 'sklearn', 'scipy', 'datetime', 'typing'],
+    verbosity_level=logging.INFO,
     planning_interval=4,
-    # prompt_templates=PROMPT_TEMPLATE["master_agent"],
+    prompt_templates=PROMPT_TEMPLATE["master_agent"],
     name="master_agent",
     description="This agent is responsible for managing audio, vlm, arithmetic and pandas agents."
 )