blasisd/GAIA_benchmarked_Agent
0
1import os2import tempfile3 4from typing import Dict, List, Optional5 6from bs4 import BeautifulSoup7import yt_dlp8import pandas as pd9import requests10import torch11 12from langchain_community.document_loaders import YoutubeLoader13from langchain_community.retrievers import BM25Retriever14from langchain_community.tools import BearlyInterpreterTool15from langchain.docstore.document import Document16from smolagents import (17 DuckDuckGoSearchTool,18 SpeechToTextTool,19 Tool,20 VisitWebpageTool,21 WikipediaSearchTool,22)23from transformers import AutoProcessor, AutoModelForImageTextToText24 25 26class RelevantInfoRetrieverTool(Tool):27 name = "relevant_info_retriever"28 description = "Retrieves relevant to the query information."29 inputs = {30 "query": {31 "type": "string",32 "description": "The query for which to retrieve information.",33 },34 "docs": {35 "type": "string",36 "description": "The source documents from which to choose in order to retrieve relevant information",37 },38 }39 output_type = "string"40 41 def forward(self, query: str, docs: List[Document]):42 self.retriever = BM25Retriever.from_documents(docs)43 results = self.retriever.get_relevant_documents(query)44 if results:45 return "\n\n".join([doc.page_content for doc in results])46 else:47 return "No relevant information found."48 49 50class YoutubeTranscriptTool(Tool):51 name = "youtube_transcript"52 description = "Fetches youtube video's transcript."53 inputs = {54 "youtube_url": {55 "type": "string",56 "description": "The youtube video url",57 },58 "source_langs": {59 "type": "array",60 "description": "A list of language codes in a descending priority for the video trascript.",61 "items": {"type": "string"},62 "default": ["en"],63 "required": False,64 "nullable": True,65 },66 "target_lang": {67 "type": "string",68 "description": "The language to which the transcript will be translated.",69 "default": "en",70 "required": False,71 "nullable": True,72 },73 }74 output_type = "string"75 76 def forward(77 self,78 youtube_url: str,79 source_langs: Optional[List[str]] = ["en"],80 target_lang: Optional[str] = "en",81 ):82 try:83 loader = YoutubeLoader.from_youtube_url(84 youtube_url,85 add_video_info=True,86 language=source_langs,87 translation=target_lang,88 # transcript_format=TranscriptFormat.CHUNKS,89 # chunk_size_seconds=30,90 )91 transcript_docs = loader.load()92 return transcript_docs93 94 except Exception as e:95 return f"Error fetching video's transcript: {e}"96 97 98class ReverseStringTool(Tool):99 name = "reverse_string"100 description = "Reverses the input string."101 inputs = {102 "string": {103 "type": "string",104 "description": "The string that needs to be reversed.",105 }106 }107 output_type = "string"108 109 def forward(self, string: str):110 try:111 return string[-1::-1]112 except Exception as e:113 return f"Error reversing string: {e}"114 115 116class SmolVLM2:117 """The parent class for visual analyzer tools (using SmolVLM2-500M-Video model)"""118 119 def __init__(self):120 """Initializations for the analyzer tool"""121 model_path = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"122 device = "cpu" # "cuda" if torch.cuda.is_available() else "cpu"123 self.processor = AutoProcessor.from_pretrained(model_path)124 self.model = AutoModelForImageTextToText.from_pretrained(125 model_path,126 torch_dtype=torch.bfloat16,127 # _attn_implementation="flash_attention_2",128 ).to(device)129 130 131class ImagesAnalyzerTool(Tool, SmolVLM2):132 name = "image_analyzer"133 description = "Analyzes each input image according to the query"134 inputs = {135 "query": {136 "type": "string",137 "description": "The query according to which the image will be analyzed.",138 },139 "images_urls": {140 "type": "array",141 "description": "A list of strings containing the images' urls",142 "items": {"type": "string"},143 },144 }145 output_type = "string"146 147 def __init__(self):148 Tool.__init__(self)149 SmolVLM2.__init__(self)150 151 def forward(self, query: str, images_urls: List[str]):152 153 try:154 155 # Image message entities for the different images' urls156 image_message_ents = [{"type": "image", "url": iu} for iu in images_urls]157 158 messages = [159 {160 "role": "user",161 "content": [162 {163 "type": "text",164 "text": query,165 },166 ]167 + image_message_ents,168 },169 ]170 171 inputs = self.processor.apply_chat_template(172 messages,173 add_generation_prompt=True,174 tokenize=True,175 return_dict=True,176 return_tensors="pt",177 ).to(self.model.device, dtype=torch.bfloat16)178 179 generated_ids = self.model.generate(180 **inputs, do_sample=False, max_new_tokens=64181 )182 generated_texts = self.processor.batch_decode(183 generated_ids,184 skip_special_tokens=True,185 )186 return generated_texts[0]187 except Exception as e:188 return f"Error analyzing image(s): {e}"189 190 191class VideoAnalyzerTool(Tool, SmolVLM2):192 name = "video_analyzer"193 description = "Analyzes video at a specified path according to the query"194 inputs = {195 "query": {196 "type": "string",197 "description": "The query according to which the video will be analyzed.",198 },199 "video_path": {200 "type": "string",201 "description": "A string containing the video path",202 },203 }204 output_type = "string"205 206 def __init__(self):207 Tool.__init__(self)208 SmolVLM2.__init__(self)209 210 def forward(self, query: str, video_path: str) -> str:211 try:212 messages = [213 {214 "role": "user",215 "content": [216 {"type": "video", "path": video_path},217 {"type": "text", "text": query},218 ],219 },220 ]221 222 inputs = self.processor.apply_chat_template(223 messages,224 add_generation_prompt=True,225 tokenize=True,226 return_dict=True,227 return_tensors="pt",228 ).to(self.model.device, dtype=torch.bfloat16)229 230 generated_ids = self.model.generate(231 **inputs, do_sample=False, max_new_tokens=64232 )233 generated_texts = self.processor.batch_decode(234 generated_ids,235 skip_special_tokens=True,236 )237 238 return generated_texts[0]239 except Exception as e:240 return f"Error analyzing video: {e}"241 finally:242 # Cleanup if needed243 if video_path and os.path.exists(video_path):244 os.remove(video_path)245 246 247class FileDownloaderTool(Tool):248 name = "file_downloader"249 description = "Downloads a file returning the name of the temporarily saved file"250 inputs = {251 "file_url": {252 "type": "string",253 "description": "The url from which the file shall be downloaded.",254 },255 }256 output_type = "string"257 258 def forward(self, file_url: str) -> str:259 response = requests.get(file_url, stream=True)260 response.raise_for_status()261 original_filename = (262 response.headers.get("content-disposition", "")263 .split("=", -1)[-1]264 .strip('"')265 )266 267 # Even if original_filename is empty or there is no extension, ext will be ""268 ext = os.path.splitext(original_filename)[-1]269 270 with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp_file:271 for chunk in response.iter_content(chunk_size=8192):272 tmp_file.write(chunk)273 return tmp_file.name274 275 276class YoutubeVideoDownloaderTool(Tool):277 name = "youtube_video_downloader"278 description = "Downloads the video from the specified url and returns the path where the video was saved"279 inputs = {280 "video_url": {281 "type": "string",282 "description": "A string containing the video url",283 },284 }285 output_type = "string"286 287 def forward(self, video_url: str) -> str:288 try:289 saved_video_path = ""290 temp_dir = tempfile.gettempdir()291 ydl_opts = {292 "outtmpl": f"{temp_dir}/%(title)s.%(ext)s", # Absolute or relative path293 "quiet": True,294 }295 296 # Download youtube video as a file in tmp directory297 with yt_dlp.YoutubeDL(ydl_opts) as ydl:298 info = ydl.extract_info(video_url, download=True)299 saved_video_path = ydl.prepare_filename(info)300 return saved_video_path301 except Exception as e:302 return f"Error downloading video: {e}"303 304 305class LoadXlsxFileTool(Tool):306 name = "load_xlsx_file"307 description = "This tool loads xlsx file into pandas and returns it"308 inputs = {"file_path": {"type": "string", "description": "File path"}}309 output_type = "object"310 311 def forward(self, file_path: str) -> object:312 return pd.read_excel(file_path)313 314 315class LoadTextFileTool(Tool):316 name = "load_text_file"317 description = "This tool loads any text file"318 inputs = {"file_path": {"type": "string", "description": "File path"}}319 output_type = "string"320 321 def forward(self, file_path: str) -> str:322 with open(file_path, "r", encoding="utf-8") as file:323 return file.read()324 325 326class WebpageTablesContextRetrieverTool(Tool):327 name = "webpage_tables_context_retriever"328 description = """Retrieves structural context for all tables on a webpage.329 Returns table indexes with captions, headers, and surrounding text to help identify relevant tables.330 Use this first to determine which table index to extract."""331 inputs = {332 "url": {"type": "string", "description": "The URL of the webpage to analyze"}333 }334 output_type = "object"335 336 def forward(self, url: str) -> Dict:337 """Retrieve context information for all tables on the page"""338 try:339 response = requests.get(url, timeout=15)340 response.raise_for_status()341 soup = BeautifulSoup(response.text, "html.parser")342 343 tables = soup.find_all("table")344 if not tables:345 return {346 "status": "success",347 "tables": [],348 "message": "No tables found on page",349 "url": url,350 }351 352 results = []353 for i, table in enumerate(tables):354 context = {355 "index": i,356 "id": table.get("id", ""),357 "class": " ".join(table.get("class", [])),358 "summary": table.get("summary", ""),359 "caption": self._get_table_caption(table),360 "preceding_header": self._get_preceding_header(table),361 "surrounding_text": self._get_surrounding_text(table),362 }363 results.append(context)364 365 return {366 "status": "success",367 "tables": results,368 "url": url,369 "message": f"Found {len(results)} tables with context information",370 "suggestion": "Use html_table_extractor with the most relevant index",371 }372 373 except Exception as e:374 return {375 "status": "error",376 "url": url,377 "message": f"Failed to retrieve table contexts: {str(e)}",378 }379 380 def _get_table_caption(self, table) -> str:381 """Extract table caption text if available"""382 caption = table.find("caption")383 return caption.get_text(strip=True) if caption else ""384 385 def _get_preceding_header(self, table) -> str:386 """Find the nearest preceding heading"""387 for tag in table.find_all_previous(["h1", "h2", "h3", "h4", "h5", "h6"]):388 return tag.get_text(strip=True)389 return ""390 391 def _get_surrounding_text(self, table, chars=150) -> str:392 """Get relevant text around the table"""393 prev_text = " ".join(394 t.strip()395 for t in table.find_all_previous(string=True, limit=3)396 if t.strip()397 )398 next_text = " ".join(399 t.strip() for t in table.find_all_next(string=True, limit=3) if t.strip()400 )401 return f"...{prev_text[-chars:]} [TABLE] {next_text[:chars]}..."402 403 404class HtmlTableExtractorTool(Tool):405 name = "html_table_extractor"406 description = """Extracts a specific HTML table as structured data.407 Use after webpage_tables_context_retriever to get the correct table index."""408 inputs = {409 "page_url": {410 "type": "string",411 "description": "The webpage URL containing the table",412 },413 "table_index": {414 "type": "integer",415 "description": "0-based index of the table to extract (from webpage_tables_context_retriever)",416 },417 }418 output_type = "object"419 420 def forward(self, page_url: str, table_index: int) -> Dict:421 """Extract a specific table by index"""422 try:423 # First verify the URL is accessible424 test_request = requests.head(page_url, timeout=5)425 test_request.raise_for_status()426 427 # Read all tables428 tables = pd.read_html(page_url)429 430 if not tables:431 return {432 "status": "error",433 "message": "No tables found at URL",434 "url": page_url,435 }436 437 # Validate index438 if table_index < 0 or table_index >= len(tables):439 return {440 "status": "error",441 "message": f"Invalid table index {table_index}. Page has {len(tables)} tables.",442 "url": page_url,443 "available_indexes": list(range(len(tables))),444 }445 446 # Convert DataFrame to JSON-serializable format447 df = tables[table_index]448 return {449 "status": "success",450 "table_index": table_index,451 "table_data": df,452 "url": page_url,453 }454 455 except Exception as e:456 return {457 "status": "error",458 "message": f"Table extraction failed: {str(e)}",459 "url": page_url,460 "table_index": table_index,461 }462