tools_smolagent_doc.py

import os
import pandas as pd
import re
import requests
import tempfile
import uuid

from smolagents import tool

from google import genai
from google.genai import types

# from smolagents import tool
from typing import Optional, Dict, Union


@tool
def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
    """
    Save content to a temporary file and return the path.
    Useful for processing files from the GAIA API.

    Args:
        content: The content to save to the file
        filename: Optional filename, will generate a random name if not provided

    Returns:
        Path to the saved file
    """
    temp_dir = tempfile.gettempdir()
    if filename is None:
        temp_file = tempfile.NamedTemporaryFile(delete=False)
        filepath = temp_file.name
    else:
        filepath = os.path.join(temp_dir, filename)

    # Write content to the file
    with open(filepath, "w") as f:
        f.write(content)

    return f"File saved to {filepath}. You can read this file to process its contents."


# File Download Tool
@tool
def download_file_from_url(
    url: str, directory: str
) -> Dict[str, Union[str, None]]:
    """Downloads a file from a URL and saves it to a directory.
    Args:
        url (str): the URL to download the file from.
        directory (str): the directory to save the file to.
    Returns:
        Dict[str, Union[str, None]]: A dictionary containing the file type and path.
    """

    try:
        response = requests.get(url, stream=True, timeout=10)
        response.raise_for_status()

        content_type = response.headers.get("content-type", "").lower()

        # Try to get filename from headers
        filename = None
        cd = response.headers.get("content-disposition", "")
        match = re.search(r"filename\*=UTF-8\'\'(.+)", cd) or re.search(
            r'filename="?([^"]+)"?', cd
        )
        if match:
            filename = match.group(1)

        # If not in headers, try URL
        if not filename:
            filename = os.path.basename(url.split("?")[0])

        # Fallback to generated filename
        if not filename:
            extension = {
                "image/jpeg": ".jpg",
                "image/png": ".png",
                "image/gif": ".gif",
                "audio/wav": ".wav",
                "audio/mpeg": ".mp3",
                "video/mp4": ".mp4",
                "text/plain": ".txt",
                "text/csv": ".csv",
                "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
                "application/vnd.ms-excel": ".xls",
                "application/octet-stream": ".bin",
            }.get(content_type, ".bin")
            filename = f"downloaded_{uuid.uuid4().hex[:8]}{extension}"

        os.makedirs(directory, exist_ok=True)
        file_path = os.path.join(directory, filename)

        with open(file_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        # shutil.copy(file_path, os.getcwd())

        if os.path.exists(file_path) and os.path.getsize(file_path) > 0:
            return {"type": content_type, "path": file_path}
        else:
            return {
                "type": "error",
                "path": None,
                "error": "Failed to save file",
            }

    except Exception as e:
        return {
            "type": "error",
            "path": None,
            "error": 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)}"


# CSV Analysis Tool
@tool
def analyze_csv_file(file_path: str, query: str) -> str:
    """Analyzes a CSV file and answers questions about its contents using Gemini.
    Args:
        file_path (str): the path to the CSV file to analyze.
        query (str): the question to answer about the CSV file.
    Returns:
        str: The result of the analysis.
    """
    try:
        # Read the CSV file
        df = pd.read_csv(file_path)

        # Initialize Gemini
        client = genai.Client(api_key=os.getenv("GEMINI_KEY"))
        model = "models/gemini-1.5-flash-8b"

        # Convert DataFrame to a string representation
        df_str = df.to_string()

        # Create a prompt for Gemini
        prompt = f"""Analyze this CSV data and provide insights:
Dimensions: {len(df)} rows × {len(df.columns)} columns
Data:
{df_str}
Please provide:
1. A summary of the data structure and content
2. Key patterns and insights
3. Potential data quality issues
4. Suggestions for analysis
User Query: {query}
Please format your response in a clear, structured way with sections and bullet points."""

        # Get analysis from Gemini
        response = client.models.generate_content(
            model=model,
            contents=types.Content(
                parts=[
                    types.Part(text=df_str),
                    types.Part(text=prompt),
                ]
            ),
        )

        result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n\n"
        result += response.text

        return result
    except Exception as e:
        return f"Error analyzing CSV file: {str(e)}"


# Excel Analysis Tool
@tool
def analyze_excel_file(file_path: str, query: str) -> str:
    """Analyzes an Excel file and answers questions about its contents using Gemini.
    Args:
        file_path (str): the path to the Excel file to analyze.
        query (str): the question to answer about the Excel file.
    Returns:
        str: The result of the analysis.
    """
    try:
        # Read all sheets from the Excel file
        excel_file = pd.ExcelFile(file_path)
        sheet_names = excel_file.sheet_names

        # Initialize Gemini
        client = genai.Client(api_key=os.getenv("GEMINI_KEY"))
        model = "models/gemini-1.5-flash-8b"

        result = f"Excel file loaded with {len(sheet_names)} sheets: {', '.join(sheet_names)}\n\n"

        # Analyze each sheet
        for sheet_name in sheet_names:
            df = pd.read_excel(file_path, sheet_name=sheet_name)

            # Convert DataFrame to a string representation
            df_str = df.to_string()

            # Create a prompt for Gemini
            prompt = f"""Analyze this Excel sheet data and provide insights:
Sheet Name: {sheet_name}
Dimensions: {len(df)} rows × {len(df.columns)} columns
Data:
{df_str}
Please provide:
1. A summary of the data structure and content
2. Key patterns and insights
3. Potential data quality issues
4. Suggestions for analysis
User Query: {query}
Please format your response in a clear, structured way with sections and bullet points."""

            # Get analysis from Gemini
            response = client.models.generate_content(
                model=model,
                contents=types.Content(
                    parts=[types.Part(text=df_str), types.Part(text=prompt)]
                ),
            )

            result += f"=== Sheet: {sheet_name} ===\n"
            result += response.text + "\n"
            result += "=" * 50 + "\n\n"

        return result
    except Exception as e:
        return f"Error analyzing Excel file: {str(e)}"


@tool
def read_file(filepath: str) -> str:
    """Reads the content of a text file.
    Args:
        filepath (str): the path to the file to read.
    Returns:
        str: The content of the file.
    """
    try:
        with open(filepath, "r", encoding="utf-8") as file:
            content = file.read()
        return content
    except FileNotFoundError:
        return f"File not found: {filepath}"
    except IOError as e:
        return f"Error reading file: {str(e)}"