punama/Final_Assignment_Template
0
1# import smolagents.models as sm_models2 3# _orig_roles = sm_models.MessageRole.roles4 5# @classmethod6# def _roles_with_control(cls):7# return _orig_roles() + ["control"]8 9# sm_models.MessageRole.roles = _roles_with_control10 11 12 13from smolagents import (CodeAgent, 14 GradioUI, 15 LiteLLMModel, 16 OpenAIServerModel, 17 ChatMessage, 18 ToolCallingAgent)19from smolagents.default_tools import (DuckDuckGoSearchTool, 20 VisitWebpageTool, 21 WikipediaSearchTool, 22 SpeechToTextTool,23 PythonInterpreterTool)24import yaml25from tools.final_answer import FinalAnswerTool, check_reasoning, ensure_formatting26from tools.tools import (youtube_frames_to_images, use_vision_model, 27 read_file, download_file_from_url, 28 extract_text_from_image, analyze_csv_file, 29 analyze_excel_file, youtube_transcribe,30 transcribe_audio, review_youtube_video,31 tavily_search)32import os33from dotenv import load_dotenv34import time35 36load_dotenv()37 38def get_gemini_keys() -> list[str]:39 keys = []40 41 multi = os.getenv("GEMINI_KEYS", "")42 if multi:43 keys.extend([k.strip() for k in multi.replace("\n", ",").split(",") if k.strip()])44 45 single = os.getenv("GEMINI_KEY", "")46 if single.strip():47 keys.append(single.strip())48 49 seen = set()50 unique = []51 for key in keys:52 if key not in seen:53 seen.add(key)54 unique.append(key)55 56 if not unique:57 raise RuntimeError("No Gemini key found. Set GEMINI_KEYS or GEMINI_KEY.")58 59 return unique60 61 62class RotatingGeminiModel:63 def __init__(64 self,65 model_id: str,66 temperature: float = 0.1,67 max_tokens: int = 512,68 sleep_seconds: int = 4,69 cooldown_seconds: int = 120,70 all_keys_wait_seconds: int = 180,71 ):72 self.model_id = model_id73 self.temperature = temperature74 self.max_tokens = max_tokens75 self.sleep_seconds = sleep_seconds76 self.cooldown_seconds = cooldown_seconds77 self.all_keys_wait_seconds = all_keys_wait_seconds78 79 self.keys = get_gemini_keys()80 self.index = 081 self.cooldowns = {i: 0 for i in range(len(self.keys))}82 83 self.models = [84 LiteLLMModel(85 model_id=model_id,86 api_key=key,87 temperature=temperature,88 max_tokens=max_tokens,89 )90 for key in self.keys91 ]92 93 def _is_rate_limit(self, error: Exception) -> bool:94 text = str(error).lower()95 return (96 "resource_exhausted" in text97 or "429" in text98 or "rate limit" in text99 or "rate_limit" in text100 or "ratelimiterror" in text101 or "generate_content_free_tier_requests" in text102 or "quota exceeded" in text103 )104 105 def __call__(self, messages, **kwargs) -> ChatMessage:106 last_error = None107 108 # Round-robin through available keys.109 for _ in range(len(self.models)):110 model_index = self.index % len(self.models)111 self.index = (self.index + 1) % len(self.models)112 113 now = time.time()114 if now < self.cooldowns[model_index]:115 continue116 117 model = self.models[model_index]118 119 try:120 time.sleep(self.sleep_seconds)121 return model(messages, **kwargs)122 123 except Exception as e:124 last_error = e125 error_text = str(e)126 127 print(f"\n--- Gemini key #{model_index + 1} failed ---")128 print(error_text[:2000])129 print("--- end Gemini error ---\n")130 131 if self._is_rate_limit(e):132 print(133 f"Gemini key #{model_index + 1} hit rate/quota limit. "134 f"Cooling down for {self.cooldown_seconds} seconds."135 )136 self.cooldowns[model_index] = time.time() + self.cooldown_seconds137 continue138 139 raise e140 141 # If all keys are cooling down, wait and retry once.142 print(f"All Gemini keys are cooling down. Waiting {self.all_keys_wait_seconds} seconds...")143 time.sleep(self.all_keys_wait_seconds)144 145 for model_index, model in enumerate(self.models):146 try:147 time.sleep(self.sleep_seconds)148 self.index = (model_index + 1) % len(self.models)149 return model(messages, **kwargs)150 151 except Exception as e:152 last_error = e153 error_text = str(e)154 155 print(f"\n--- Gemini key #{model_index + 1} failed after all-keys wait ---")156 print(error_text[:2000])157 print("--- end Gemini retry error ---\n")158 159 if self._is_rate_limit(e):160 self.cooldowns[model_index] = time.time() + self.cooldown_seconds161 continue162 163 raise e164 165 raise last_error166 167# Load prompts from YAML file168with open("prompts.yaml", 'r') as stream:169 prompt_templates = yaml.safe_load(stream)170 171 172# class ThinkingLiteLLMModel(LiteLLMModel):173# def __init__(self, *args, **kwargs):174# # ensure the Litellm client also maps "control" → "control"175# cr = kwargs.pop("custom_role_conversions", {})176# cr["control"] = "control"177# super().__init__(*args, custom_role_conversions=cr, **kwargs)178 179# def __call__(self, messages, **kwargs) -> ChatMessage:180 # NOTE: content must be a list of {type, text} dicts181# thinking_msg = {182# "role": "control",183# "content": [{"type": "text", "text": "thinking"}]184# }185# # prepend onto whatever messages the Agent built186# return super().__call__([thinking_msg] + messages, **kwargs)187 188class SlowLiteLLMModel(LiteLLMModel):189 def __init__(self, *args, **kwargs):190 super().__init__(*args, **kwargs)191 192 def __call__(self, messages, **kwargs) -> ChatMessage:193 time.sleep(20)194 # prepend onto whatever messages the Agent built195 return super().__call__(messages, **kwargs)196 197# # search_model_name = 'granite3.3:latest'198# search_model_name = 'cogito:14b'199# # search_model_name = 'qwen2:7b'200# search_model = ThinkingLiteLLMModel(model_id=f'ollama_chat/{search_model_name}',201# flatten_messages_as_text=True)202 203# web_agent = CodeAgent(204# model=search_model,205# tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), FinalAnswerTool()],206# max_steps=6,207# verbosity_level=1,208# grammar=None,209# planning_interval=4,210# name="web_agent",211# description="Searches the web using the and reviews web pages to find information.",212# additional_authorized_imports=['bs4', 'requests', 'io', 'wiki'],213# prompt_templates=prompt_templates214# )215 216# image_model_name = 'llama3.2-vision'217# image_model = OpenAIServerModel(model_id=image_model_name,218# api_base='http://localhost:11434/v1/',219# api_key='ollama',220# flatten_messages_as_text=False)221# image_agent = ToolCallingAgent(222# model=image_model,223# tools=[FinalAnswerTool()],224# max_steps=4,225# verbosity_level=2,226# grammar=None,227# planning_interval=4,228# #additional_authorized_imports=["PIL", "requests", "io", "numpy"],229# name="image_agent",230# description="Review images and videos for answers to questions based on visual data",231# prompt_templates=prompt_templates232# )233 234# react_model_name = 'qwen2:7b'235# # Initialize the chat model236# react_model = OpenAIServerModel(model_id=react_model_name,237# api_base='http://localhost:11434/v1/',238# api_key='ollama',239# flatten_messages_as_text=False)240 241react_model_name = "gemini/gemini-2.5-flash"242 243react_model = RotatingGeminiModel(244 model_id=react_model_name,245 temperature=0.1,246 max_tokens=512,247 sleep_seconds=4,248 cooldown_seconds=120,249 all_keys_wait_seconds=180,250)251 252 253manager_agent = CodeAgent(254 model=react_model,255 tools=[FinalAnswerTool(), 256 DuckDuckGoSearchTool(), 257 VisitWebpageTool(max_output_length=500000), 258 WikipediaSearchTool(extract_format='HTML'),259 tavily_search,260 SpeechToTextTool(),261 youtube_frames_to_images,262 youtube_transcribe,263 use_vision_model,264 read_file, download_file_from_url, 265 extract_text_from_image, 266 analyze_csv_file, analyze_excel_file,267 transcribe_audio,268 review_youtube_video269 ],270 managed_agents=[],271 additional_authorized_imports=['os', 'pandas', 'numpy', 'PIL', 'tempfile', 'PIL.Image'],272 max_steps=12,273 verbosity_level=1,274 planning_interval=4,275 name="Manager",276 description="The manager of the team, responsible for overseeing and guiding the team's work.",277 final_answer_checks=[],278 prompt_templates=prompt_templates279)280 281 282 283if __name__ == "__main__":284 GradioUI(manager_agent).launch()285