Create tools.py

tools.py

@@ -0,0 +1,431 @@
+import tempfile
+import requests
+import os
+
+from time import sleep
+from urllib.parse import urlparse
+from typing import Optional, List
+import yt_dlp
+import imageio
+from google.genai import types
+
+from PIL import Image
+from smolagents import CodeAgent, tool, OpenAIServerModel, LiteLLMModel
+from google import genai
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> None:
+#     sleep(1.0)  # Let JavaScript animations happen before taking the screenshot
+#     driver = helium.get_driver()
+#     current_step = memory_step.step_number
+#     if driver is not None:
+#         for previous_memory_step in agent.memory.steps:  # Remove previous screenshots from logs for lean processing
+#             if isinstance(previous_memory_step, ActionStep) and previous_memory_step.step_number <= current_step - 2:
+#                 previous_memory_step.observations_images = None
+#         png_bytes = driver.get_screenshot_as_png()
+#         image = Image.open(BytesIO(png_bytes))
+#         print(f"Captured a browser screenshot: {image.size} pixels")
+#         memory_step.observations_images = [image.copy()]  # Create a copy to ensure it persists, important!
+
+#     # Update observations with current URL
+#     url_info = f"Current url: {driver.current_url}"
+#     memory_step.observations = (
+#         url_info if memory_step.observations is None else memory_step.observations + "\n" + url_info
+#     )
+#     return
+
+# def initialize_driver():
+#     """Initialize the Selenium WebDriver."""
+#     chrome_options = webdriver.ChromeOptions()
+#     chrome_options.add_argument("--force-device-scale-factor=1")
+#     chrome_options.add_argument("--window-size=1000,1350")
+#     chrome_options.add_argument("--disable-pdf-viewer")
+#     chrome_options.add_argument("--window-position=0,0")
+#     return helium.start_chrome(headless=False, options=chrome_options)
+
+# @tool
+# def search_item_ctrl_f(text: str, nth_result: int = 1) -> str:
+#     """
+#     Searches for text on the current page via Ctrl + F and jumps to the nth occurrence.
+#     Args:
+#         text: The text to search for
+#         nth_result: Which occurrence to jump to (default: 1)
+#     """
+#     elements = driver.find_elements(By.XPATH, f"//*[contains(text(), '{text}')]")
+#     if nth_result > len(elements):
+#         raise Exception(f"Match n°{nth_result} not found (only {len(elements)} matches found)")
+#     result = f"Found {len(elements)} matches for '{text}'."
+#     elem = elements[nth_result - 1]
+#     driver.execute_script("arguments[0].scrollIntoView(true);", elem)
+#     result += f"Focused on element {nth_result} of {len(elements)}"
+#     return result
+
+
+# @tool
+# def go_back() -> None:
+#     """Used when navigating web pages using Helium.  Goes back to previous page."""
+#     driver.back()
+
+
+# @tool
+# def close_popups() -> str:
+#     """
+#     Closes any visible modal or pop-up on the page. Use this to dismiss pop-up windows! This does not work on cookie consent banners.
+#     """
+#     webdriver.ActionChains(driver).send_keys(Keys.ESCAPE).perform()
+
+@tool
+def use_vision_model(question: str, images: List[Image.Image]) -> str:
+    """
+    Use a Vision Model to answer a question about a set of images.  
+    Always use this tool to ask questions about a set of images you have been provided.
+    This function uses an image-to-text AI model.  
+    You can ask a question about a list of one image or a list of multiple images.  
+    So, if you have multiple images that you want to ask the same question of, pass the entire list of images to the model.
+    Ensure your prompt is specific enough to retrieve the exact information you are looking for.
+    
+    Args:
+        question: The question to ask about the images.  Type: str
+        images: The list of images to as the question about.  Type: List[PIL.Image.Image]
+    """
+    image_model_name = "gemini/gemini-1.5-flash"
+    
+    print(f'Leveraging model {image_model_name}')
+    # image_model_name = 'gemma3:12b'
+    # image_model = OpenAIServerModel(
+    #     model_id=image_model_name,
+    #     api_base='http://localhost:11434/v1/',
+    #     api_key='ollama',
+    #     flatten_messages_as_text=False
+    # )
+    image_model =LiteLLMModel(model_id=image_model_name, 
+                           api_key=os.getenv("GEMINI_KEY"),
+                           temperature=0.2
+                           )
+
+    content = [
+        {
+            "type": "text",
+            "text": question
+        }
+    ]
+    print(f"Asking model a question about {len(images)} images")
+    for image in images:
+        content.append({
+            "type": "image",
+            "image": image  # ✅ Directly the PIL Image, no wrapping
+        })
+
+    messages = [
+        {
+            "role": "user",
+            "content": content
+        }
+    ]
+
+    output = image_model(messages).content
+    print(f'Model returned: {output}')
+    return output
+
+@tool
+def review_youtube_video(url: str, question: str) -> str:
+    """
+    Reviews a YouTube video and answers a specific question about that video.  
+
+    Args:
+        url (str): the URL to the YouTube video.  Should be like this format: https://www.youtube.com/watch?v=9hE5-98ZeCg
+        question (str): The question you are asking about the video
+    """
+    try:
+        client = genai.Client(api_key=os.getenv('GEMINI_KEY'))
+        model = 'gemini-2.0-flash-lite'
+        response = client.models.generate_content(
+            model=model,
+            contents=types.Content(
+                parts=[
+                    types.Part(
+                        file_data=types.FileData(file_uri=url)
+                    ),
+                    types.Part(text=question)
+                ]
+            )
+        )
+        return response.text
+    except Exception as e:
+        return f"Error asking {model} about video: {str(e)}"
+
+@tool
+def youtube_frames_to_images(url: str, sample_interval_seconds: int = 5) -> List[Image.Image]:
+    """
+    Reviews a YouTube video and returns a List of PIL Images (List[PIL.Image.Image]), which can then be reviewed by a vision model.
+    Only use this tool if you have been given a YouTube video that you need to analyze.
+    This will generate a list of images, and you can use the use_vision_model tool to analyze those images
+    Args:
+        url: The Youtube URL
+        sample_interval_seconds: The sampling interval (default is 5 seconds)
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Download the video locally
+        ydl_opts = {
+            'format': 'bestvideo[height<=1080]+bestaudio/best[height<=1080]/best',
+            'outtmpl': os.path.join(tmpdir, 'video.%(ext)s'),
+            'quiet': True,
+            'noplaylist': True,
+            'merge_output_format': 'mp4',
+            'force_ipv4': True,  # Avoid IPv6 issues
+        }
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info = ydl.extract_info(url, download=True)
+        
+        # Find the downloaded file
+        video_path = None
+        for file in os.listdir(tmpdir):
+            if file.endswith('.mp4'):
+                video_path = os.path.join(tmpdir, file)
+                break
+        
+        if not video_path:
+            raise RuntimeError("Failed to download video as mp4")
+
+        # ✅ Fix: Use `imageio.get_reader()` instead of `imopen()`
+        reader = imageio.get_reader(video_path)  # Works for frame-by-frame iteration
+        metadata = reader.get_meta_data()
+        fps = metadata.get('fps')
+        
+        if fps is None:
+            reader.close()
+            raise RuntimeError("Unable to determine FPS from video metadata")
+
+        frame_interval = int(fps * sample_interval_seconds)
+        images: List[Image.Image] = []
+
+        # ✅ Iterate over frames using `get_reader()`
+        for idx, frame in enumerate(reader):
+            if idx % frame_interval == 0:
+                images.append(Image.fromarray(frame))
+
+        reader.close()
+        return images
+
+@tool
+def read_file(filepath: str ) -> str:
+    """
+    Used to read the content of a file.  Returns the content as a string.
+    Will only work for text-based files, such as .txt files or code files.
+    Do not use for audio or visual files. 
+    
+    Args:
+        filepath (str): The path to the file to be read.
+
+    Returns:
+        str: Content of the file as a string.
+    
+    Raises:
+        IOError: If there is an error opening or reading from the file.
+    """
+    try:
+        with open(filepath, 'r', encoding='utf-8') as file:
+            content = file.read()
+        print(content)
+        return content
+    except FileNotFoundError:
+        print(f"File not found: {filepath}")
+    except IOError as e:
+        print(f"Error reading file: {str(e)}")
+
+@tool
+def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
+    """
+    Download a file from a URL and save it to a temporary location.  
+    Use this tool when you are asked a question and told that there is a file or image provided.
+
+    
+    Args:
+        url: The URL to download from
+        filename: Optional filename, will generate one based on URL if not provided
+        
+    Returns:
+        Path to the downloaded file
+    """
+    try:
+        # Parse URL to get filename if not provided
+        print(f"Downloading file from {url}")
+        if not filename:
+            path = urlparse(url).path
+            filename = os.path.basename(path)
+            if not filename:
+                # Generate a random name if we couldn't extract one
+                import uuid
+                filename = f"downloaded_{uuid.uuid4().hex[:8]}"
+        
+        # Create temporary file
+        temp_dir = tempfile.gettempdir()
+        filepath = os.path.join(temp_dir, filename)
+        
+        # Download the file
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        
+        # Save the file
+        with open(filepath, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        
+        return f"File downloaded to {filepath}. You can now process this file."
+    except Exception as e:
+        return f"Error downloading file: {str(e)}"
+
+@tool
+def extract_text_from_image(image_path: str) -> str:
+    """
+    Extract text from an image using pytesseract (if available).
+    
+    Args:
+        image_path: Path to the image file
+        
+    Returns:
+        Extracted text or error message
+    """
+    try:
+        # Try to import pytesseract
+        import pytesseract
+        from PIL import Image
+        
+        # Open the image
+        image = Image.open(image_path)
+        
+        # Extract text
+        text = pytesseract.image_to_string(image)
+        
+        return f"Extracted text from image:\n\n{text}"
+    except ImportError:
+        return "Error: pytesseract is not installed. Please install it with 'pip install pytesseract' and ensure Tesseract OCR is installed on your system."
+    except Exception as e:
+        return f"Error extracting text from image: {str(e)}"
+
+@tool
+def analyze_csv_file(file_path: str, query: str) -> str:
+    """
+    Analyze a CSV file using pandas and answer a question about it.  
+    To use this file you need to have saved it in a location and pass that location to the function.
+    The download_file_from_url tool will save it by name to tempfile.gettempdir()
+    
+    Args:
+        file_path: Path to the CSV file
+        query: Question about the data
+        
+    Returns:
+        Analysis result or error message
+    """
+    try:
+        import pandas as pd
+        
+        # Read the CSV file
+        df = pd.read_csv(file_path)
+        
+        # Run various analyses based on the query
+        result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
+        result += f"Columns: {', '.join(df.columns)}\n\n"
+        
+        # Add summary statistics
+        result += "Summary statistics:\n"
+        result += str(df.describe())
+        
+        return result
+    except ImportError:
+        return "Error: pandas is not installed. Please install it with 'pip install pandas'."
+    except Exception as e:
+        return f"Error analyzing CSV file: {str(e)}"
+
+@tool
+def analyze_excel_file(file_path: str, query: str) -> str:
+    """
+    Analyze an Excel file using pandas and answer a question about it.
+    To use this file you need to have saved it in a location and pass that location to the function.
+    The download_file_from_url tool will save it by name to tempfile.gettempdir()
+    
+    Args:
+        file_path: Path to the Excel file
+        query: Question about the data
+        
+    Returns:
+        Analysis result or error message
+    """
+    try:
+        import pandas as pd
+        
+        # Read the Excel file
+        df = pd.read_excel(file_path)
+        
+        # Run various analyses based on the query
+        result = f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
+        result += f"Columns: {', '.join(df.columns)}\n\n"
+        
+        # Add summary statistics
+        result += "Summary statistics:\n"
+        result += str(df.describe())
+        
+        return result
+    except ImportError:
+        return "Error: pandas and openpyxl are not installed. Please install them with 'pip install pandas openpyxl'."
+    except Exception as e:
+        return f"Error analyzing Excel file: {str(e)}"
+
+import whisper
+
+@tool
+def youtube_transcribe(url: str) -> str:
+    """
+    Transcribes a YouTube video.  Use when you need to process the audio from a YouTube video into Text.
+
+    Args:
+        url: Url of the YouTube video
+    """
+    model_size: str = "small"
+    # Load model
+    model = whisper.load_model(model_size)
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Download audio
+        ydl_opts = {
+            'format': 'bestaudio/best',
+            'outtmpl': os.path.join(tmpdir, 'audio.%(ext)s'),
+            'quiet': True,
+            'noplaylist': True,
+            'postprocessors': [{
+                'key': 'FFmpegExtractAudio',
+                'preferredcodec': 'wav',
+                'preferredquality': '192',
+            }],
+            'force_ipv4': True,
+        }
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info = ydl.extract_info(url, download=True)
+
+        audio_path = next((os.path.join(tmpdir, f) for f in os.listdir(tmpdir) if f.endswith('.wav')), None)
+        if not audio_path:
+            raise RuntimeError("Failed to find audio")
+
+        # Transcribe
+        result = model.transcribe(audio_path)
+        return result['text']
+
+@tool
+def transcribe_audio(audio_file_path: str) -> str:
+    """
+    Transcribes an audio file.  Use when you need to process audio data.
+    DO NOT use this tool for YouTube video; use the youtube_transcribe tool to process audio data from YouTube.
+    Use this tool when you have an audio file in .mp3, .wav, .aac, .ogg, .flac, .m4a, .alac or .wma
+
+    Args:
+        audio_file_path: Filepath to the audio file (str)
+    """
+    model_size: str = "small"
+    # Load model
+    model = whisper.load_model(model_size)
+    result = model.transcribe(audio_file_path)
+    return result['text']
+
+# global driver
+# driver = initialize_driver()