agent_tools.py

import base64
import mimetypes
import os
from pathlib import Path
import subprocess
import sys
from urllib.parse import urlparse
import re

from bs4 import BeautifulSoup
from ddgs import DDGS
from langchain_core.messages import HumanMessage
from langchain_google_genai import ChatGoogleGenerativeAI
import pandas as pd
import requests
from langchain_core.tools import tool

from logging_config import get_logger

logger = get_logger(__name__)
tools_llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
DEFAULT_DOWNLOAD_DIR = "/tmp/agent_files"


@tool
def web_search(query: str) -> str:
    """
    Tool name: web_search
    Description: Use this tool when the user asks for a summary of web search results about a topic
    the query param should be something very simple and short.
    Input: A string containing the search query (e.g., 'latest AI research trends in 2025')
    Output: A string containing a formatted response in this structure:

    Title:
    <Extracted Title>

    Body:
    <Summary of the findings>

    Reference:
    <URL or source name>

    This tool searches the web using DuckDuckGo, stores the results in DuckDB, filters and summarizes them.
    Use only when the user explicitly asks for updated, online, or news-related information.
    """
    raw_search_results = _search_duckduckgo(query)
    enriched_search_results = _enrich_web_search_results(raw_search_results)
    formatted_search_results = _format_web_search_output(enriched_search_results)
    return formatted_search_results

@tool
def download_file(url: str) -> dict:
    """Download a file from a URL.

    Args:
        url: The URL of the file to download.

    Returns:
        dict: A dictionary containing the path to the downloaded file and a dictionary with metadata about the file.
    """
    logger.info(f"Downloading file from {url}")
    response = requests.get(url, stream=True)
    response.raise_for_status()

    parsed = urlparse(url)
    base = os.path.basename(parsed.path)
    file_name, file_extension = os.path.splitext(base)
    file_extension = file_extension.lower()

    # Si no hay extensión en URL, intentar con Content-Disposition
    if not file_extension:
        cd = response.headers.get('content-disposition', '')
        if cd:
            match = re.search(r'filename\*?=(?:UTF-8\'\')?"?([^\";]+)"?', cd)
            if match:
                fname = os.path.basename(match.group(1))
                name2, ext2 = os.path.splitext(fname)
                if ext2:
                    file_name, file_extension = name2, ext2.lower()

    # Si aún sin extensión, dejar ext vacía
    if not file_name:
        file_name = "downloaded_file"


    filename = f"{file_name}{file_extension}"
    full_path = os.path.join(DEFAULT_DOWNLOAD_DIR, filename)

    size = 0
    os.makedirs(DEFAULT_DOWNLOAD_DIR, exist_ok=True)
    with open(full_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            if chunk:
                size += len(chunk)
                f.write(chunk)

    metadata = {"file_name": file_name,
            "file_extension": file_extension,
            "size_bytes": size,
            "file_path": os.path.abspath(full_path)}

    if file_extension in (".csv", ".xls", ".xlsx", ".xlsm", ".xlsb", ".ods"):
        try:
            df = pd.read_excel(full_path) if file_extension != ".csv" else pd.read_csv(full_path)
            metadata["type"] = "table"
            metadata["num_rows"], metadata["num_columns"] = df.shape
            metadata["columns"] = [
                {"name": col, "dtype": str(df[col].dtype)} for col in df.columns
            ]
            if df.shape[0] >= 1:
                first_row = df.iloc[0].to_dict()
                metadata["first_row"] = first_row
        except Exception:
            pass

    logger.info(f"File downloaded: {os.path.abspath(full_path)}")
    return {"file_path": os.path.abspath(full_path), "metadata": metadata}


@tool
def query_spreadsheet(file_path: str, pandas_query: str) -> str:
    """Execute a pandas query on a spreadsheet file.

    Args:
        file_path: The path to the spreadsheet file.
        pandas_query: The pandas code to execute.

    Returns:
        str: The result of the pandas code execution.
    """
    logger.info(f"Querying spreadsheet: {file_path} with code: {pandas_query}")
    if file_path.endswith(".csv"):
        df = pd.read_csv(file_path)
    elif file_path.endswith((".xls", ".xlsx")):
        df = pd.read_excel(file_path)
    else:
        raise ValueError("Formato no soportado")

    # Ejecutar el código generado por el LLM
    local_vars = {"df": df}
    try:
        exec("result = " + pandas_query, {}, local_vars)
        result = local_vars["result"]
        logger.info(f"Spreadsheet query result: {result}")
        return result.to_string(index=False) if hasattr(result, "to_string") else str(result)
    except Exception as e:
        logger.error(f"Error executing pandas query: {e}")
        return f"Error executing pandas query: {e}"


@tool
def query_media_file(file_path: str, query: str) -> str:
    """Query a media file (image or audio) for information.
    
    Args:
        file_path: Path to the image or audio file
        query: The query asking about information in the file. Be as specific as possible with the query.

    Returns:
        str: A string with the answer to the query.
    """
    logger.info(f"Reading media file: {file_path}")
    logger.info(f"Querying media file: {query}")
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File not found: {file_path}")
    
    # Get MIME type to determine if it's image or audio
    mime_type, _ = mimetypes.guess_type(file_path)
    

    if mime_type and mime_type.startswith('image/'):
        message = HumanMessage(
        content= [
            {"type": "text", "text": query},
            _encode_image(file_path)
            ]
        )
    elif mime_type and mime_type.startswith('audio/'):
        message = HumanMessage(
        content= [
            {"type": "text", "text": query},
            _encode_audio(file_path)
            ]
        )
    else:
        raise ValueError(f"Unsupported file type: {mime_type}")
    
    result = tools_llm.invoke([message])
    logger.info(f"🤖 Read media file tool response: {result.content}")
    return result.content

@tool
def execute_python_code(file_path: str) -> str:
    """Execute a Python code file.

    Args:
        file_path: The path to the Python code file.

    Returns:
        str: The result of the Python code execution.
    """
    logger.info(f"Executing Python code from: {file_path}")
    try:
        result = subprocess.run(
            [sys.executable, file_path], 
            capture_output=True, 
            text=True, 
            timeout=30  # Prevent hanging
        )
        if result.returncode == 0:
            logger.info(f"Python code executed successfully. Result: {result.stdout}")
            return result.stdout
        else:
            logger.error(f"Python code execution failed. Error: {result.stderr}")
            return f"Error: {result.stderr}"
    except subprocess.TimeoutExpired:
        logger.error("Python code execution timed out")
        return "Error: Script execution timed out"
    except Exception as e:
        logger.error(f"Error executing script: {str(e)}")
        return f"Error executing script: {str(e)}"

def _search_duckduckgo(query: str) -> list[dict[str, str]]:
    """Performs a web search using DuckDuckGo.

    Args:
        query: A string containing the search term.

    Returns:
        A list of web search results stored in dictionaries with 'title', 'href', 'body'
    """
    logger.info("🔍 Starting DuckDuckGo search with query: '%s'", query)
    results = []

    with DDGS() as ddgs:
        for r in ddgs.text(query, max_results=5):
            results.append(
                {"title": r["title"], "href": r["href"], "body": r.get("body", "")}
            )

    logger.info("✅ DuckDuckGo search completed. Found %d results.", len(results))
    logger.info(f"🔍 DuckDuckGo search results: {results}")
    return results


def _enrich_web_search_results(search_results: list[dict[str, str]]) -> list[dict[str, str]]:
    """Enhances the search result bodies by scraping full page text from each URL.

    Args:
        search_results: A list of dictionaries with 'href'

    Returns:
        A list of enriched web search results stored in dictionaries with 'title', 'href', 'body'
    """
    logger.info("🌐 Enriching search result bodies with full web content.")
    enriched_results = []

    for result in search_results:
        url = result["href"]
        try:
            logger.info(f"🔗 Fetching content from: {url}")
            headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}
            response = requests.get(url, headers=headers, timeout=5)

            if response.status_code == 200:
                soup = BeautifulSoup(response.text, "html.parser")
                # Get main textual content
                texts = soup.stripped_strings
                full_text = " ".join(texts)
                # Truncate for safety (optional)
                result["body"] = full_text[:3000]
                logger.info("✅ Content fetched successfully.")
            else:
                logger.info(f"⚠️ Failed to fetch content. Status code: {response.status_code}")
        except Exception as e:
            logger.info(f"⚠️ Error scraping {url}: {str(e)}")
            # keep original body
        enriched_results.append(result)

    logger.info(f"🌐 Enrichment completed. {enriched_results}")
    return enriched_results


def _format_web_search_output(search_results: list[dict[str, str]]) -> str:
    """Formats the search result output into a readable string.

    Args:
        search_results: A list of dictionaries with 'title', 'href', 'body'

    Returns:
        A string containing a formatted response in this structure:

    Title:
    <Extracted Title>

    Body:
    <Summary of the findings>

    Reference:
    <URL or source name>
    """
    logger.info("📝 Formatting search results output.")
    if not search_results:
        response = "No relevant results were found for your search."
        logger.info("❌ No relevant results were found for your search.")
    else:
        lines = [
            f"- Title: {search_result['title']} \n Body: {search_result['body']} \n  Reference: ({search_result['href']}) \n"
            for search_result in search_results
        ]
        response = "Here are some relevant results:\n" + "\n".join(lines)


    logger.info(f"✅ Output formatted. {response}")
    return response

def _encode_image(image_path: str) -> dict:
    """Encode an image file to base64 format for Gemini model. Supports: PNG, JPEG, WEBP, HEIC, HEIF"""
    logger.info(f"Encoding image file: {image_path}")
    path = Path(image_path)
    if not path.exists():
        raise FileNotFoundError(f"Image file not found: {image_path}")
    
    # Get MIME type
    mime_type, _ = mimetypes.guess_type(image_path)
    if not mime_type or not mime_type.startswith('image/'):
        raise ValueError(f"Unsupported image format: {mime_type}")
    
    # Read and encode image
    with open(image_path, "rb") as image_file:
        encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
    
    logger.info(f"Image encoded: {image_path}")
    return {
        "type": "image_url",
        "image_url": f"data:image/png;base64,{encoded_image}"
    }

def _encode_audio(audio_path: str) -> dict:
    """Encode an audio file to base64 format for Gemini model. Supports: MP3, MPEG, MP4, MPG, AVI, WMV, MPEGPS, FLV"""
    logger.info(f"Encoding audio file: {audio_path}")
    path = Path(audio_path)
    if not path.exists():
        raise FileNotFoundError(f"Audio file not found: {audio_path}")
    
    # Get MIME type
    mime_type, _ = mimetypes.guess_type(audio_path)
    if not mime_type or not mime_type.startswith('audio/'):
        # Handle common audio formats that might not be detected
        if audio_path.lower().endswith('.mp3'):
            mime_type = 'audio/mpeg'
        elif audio_path.lower().endswith('.wav'):
            mime_type = 'audio/wav'
        elif audio_path.lower().endswith('.m4a'):
            mime_type = 'audio/mp4'
        else:
            raise ValueError(f"Unsupported audio format: {audio_path}")
    
    # Read and encode audio
    with open(audio_path, "rb") as audio_file:
        encoded_string = base64.b64encode(audio_file.read()).decode('utf-8')
    
    logger.info(f"Audio encoded: {audio_path}")
    return {
        "type": "media",
        "mime_type": mime_type,
        "data": encoded_string
    }