Update agent.py

agent.py

@@ -1,36 +1,148 @@
+# import smolagents.models as sm_models
+
+# _orig_roles = sm_models.MessageRole.roles
+
+# @classmethod
+# def _roles_with_control(cls):
+#     return _orig_roles() + ["control"]
+
+# sm_models.MessageRole.roles = _roles_with_control
+
+
+
+from smolagents import (CodeAgent, 
+                        GradioUI, 
+                        LiteLLMModel, 
+                        OpenAIServerModel, 
+                        ChatMessage, 
+                        ToolCallingAgent)
+from smolagents.default_tools import (DuckDuckGoSearchTool, 
+                                      VisitWebpageTool, 
+                                      WikipediaSearchTool, 
+                                      SpeechToTextTool,
+                                      PythonInterpreterTool)
+import yaml
+from tools.final_answer import FinalAnswerTool, check_reasoning, ensure_formatting
+from tools.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)
 import os
-import requests
+from dotenv import load_dotenv
+import time
+
+load_dotenv()
+
+# Load prompts from YAML file
+with open("prompts.yaml", 'r') as stream:
+    prompt_templates = yaml.safe_load(stream)
+
+
+# class ThinkingLiteLLMModel(LiteLLMModel):
+#     def __init__(self, *args, **kwargs):
+#         # ensure the Litellm client also maps "control" → "control"
+#         cr = kwargs.pop("custom_role_conversions", {})
+#         cr["control"] = "control"
+#         super().__init__(*args, custom_role_conversions=cr, **kwargs)
 
-API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-alpha"
-HEADERS = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
+#     def __call__(self, messages, **kwargs) -> ChatMessage:
+        # NOTE: content must be a list of {type, text} dicts
+#         thinking_msg = {
+#             "role": "control",
+#             "content": [{"type": "text", "text": "thinking"}]
+#         }
+#         # prepend onto whatever messages the Agent built
+#         return super().__call__([thinking_msg] + messages, **kwargs)
+    
+class SlowLiteLLMModel(LiteLLMModel):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
 
-system_prompt = (
-    "You are an expert AI assistant taking a benchmark test for reasoning. "
-    "Answer only the main question directly. Do not explain. Do not show work. "
-    "If you see images or documents mentioned, assume you have access and extract answer. "
-    "Use search or tools if needed. Answer format must be plain with no prefix or suffix."
+    def __call__(self, messages, **kwargs) -> ChatMessage:
+        time.sleep(15)
+        # prepend onto whatever messages the Agent built
+        return super().__call__(messages, **kwargs)
+
+# # search_model_name = 'granite3.3:latest'
+# search_model_name = 'cogito:14b'
+# # search_model_name = 'qwen2:7b'
+# search_model = ThinkingLiteLLMModel(model_id=f'ollama_chat/{search_model_name}',
+#                              flatten_messages_as_text=True)
+
+# web_agent = CodeAgent(
+#     model=search_model,
+#     tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), FinalAnswerTool()],
+#     max_steps=6,
+#     verbosity_level=1,
+#     grammar=None,
+#     planning_interval=6,
+#     name="web_agent",
+#     description="Searches the web using the and reviews web pages to find information.",
+#     additional_authorized_imports=['bs4', 'requests', 'io', 'wiki'],
+#     prompt_templates=prompt_templates
+# )
+
+# image_model_name = 'llama3.2-vision'
+# image_model = OpenAIServerModel(model_id=image_model_name,
+#                                 api_base='http://localhost:11434/v1/',
+#                                 api_key='ollama',
+#                             flatten_messages_as_text=False)
+# image_agent = ToolCallingAgent(
+#     model=image_model,
+#     tools=[FinalAnswerTool()],
+#     max_steps=4,
+#     verbosity_level=2,
+#     grammar=None,
+#     planning_interval=6,
+#     #additional_authorized_imports=["PIL", "requests", "io", "numpy"],
+#     name="image_agent",
+#     description="Review images and videos for answers to questions based on visual data",
+#     prompt_templates=prompt_templates
+# )
+
+# react_model_name = 'qwen2:7b'
+# # Initialize the chat model
+# react_model = OpenAIServerModel(model_id=react_model_name,
+#                                 api_base='http://localhost:11434/v1/',
+#                                 api_key='ollama',
+#                             flatten_messages_as_text=False)
+
+react_model_name = "gemini/gemini-2.5-flash-preview-04-17"
+react_model = LiteLLMModel(model_id=react_model_name, 
+                           api_key=os.getenv("GEMINI_KEY"),
+                           temperature=0.2
+                           )
+
+
+manager_agent = CodeAgent(
+    model=react_model,
+    tools=[FinalAnswerTool(), 
+           DuckDuckGoSearchTool(), 
+           VisitWebpageTool(max_output_length=500000), 
+           WikipediaSearchTool(extract_format='HTML'),
+           SpeechToTextTool(),
+           youtube_frames_to_images,
+           youtube_transcribe,
+           use_vision_model,
+           read_file, download_file_from_url, 
+           extract_text_from_image, 
+           analyze_csv_file, analyze_excel_file,
+           transcribe_audio,
+           review_youtube_video
+           ],
+    managed_agents=[],
+    additional_authorized_imports=['os', 'pandas', 'numpy', 'PIL', 'tempfile', 'PIL.Image'],
+    max_steps=20,
+    verbosity_level=1,
+    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],
+    prompt_templates=prompt_templates
 )
 
-class BasicAgent:
-    def __init__(self):
-        print("✅ Agent initialized with Zephyr 7B.")
-
-    def __call__(self, question: str) -> str:
-        prompt = f"{system_prompt}\n\nQuestion: {question}\nAnswer:"
-        try:
-            response = requests.post(
-                API_URL,
-                headers=HEADERS,
-                json={"inputs": prompt},
-                timeout=40
-            )
-            response.raise_for_status()
-            result = response.json()
-            if isinstance(result, list):
-                return result[0]["generated_text"].split("Answer:")[-1].strip()
-            elif "generated_text" in result:
-                return result["generated_text"].strip()
-            else:
-                return "ERROR: Unexpected response format"
-        except Exception as e:
-            return f"AGENT ERROR: {e}"
+
+
+if __name__ == "__main__":
+    GradioUI(manager_agent).launch()