tools.py

import base64
import cmath
import os
import re
import subprocess
import sys
import tempfile
import uuid
from io import BytesIO
from typing import Any, Dict, List, Optional
from urllib.parse import urlparse

import numpy as np
import pandas as pd
import requests
from dotenv import load_dotenv
from langchain.tools import tool
from langchain_community.document_loaders import (
    ArxivLoader,
    PyMuPDFLoader,
    TextLoader,
    UnstructuredExcelLoader,
    UnstructuredPDFLoader,
    WikipediaLoader,
    YoutubeLoader,
)
from langchain_community.document_loaders.csv_loader import UnstructuredCSVLoader
from langchain_tavily import TavilySearch
from PIL import Image, ImageDraw, ImageEnhance, ImageFilter, ImageFont

load_dotenv()


# ── Image helpers ─────────────────────────────────────────────────────────────


def decode_image(image_base64: str) -> Image.Image:
    data = base64.b64decode(image_base64)
    return Image.open(BytesIO(data))


def encode_image(path: str) -> str:
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def save_image(img: Image.Image, subdir: str = None) -> str:
    temp_dir = tempfile.gettempdir()
    if subdir:
        temp_dir = os.path.join(temp_dir, subdir)
        os.makedirs(temp_dir, exist_ok=True)
    filepath = os.path.join(temp_dir, f"{uuid.uuid4().hex[:8]}.png")
    img.save(filepath)
    return filepath


# ── Code interpreter ──────────────────────────────────────────────────────────


class CodeInterpreter:
    """Minimal interpreter: supports Python and Bash execution via subprocess."""

    def execute_code(self, code: str, language: str = "python") -> dict:
        if language == "python":
            return self._run([sys.executable, "-c", code])
        if language == "bash":
            return self._run(code, shell=True)
        return {
            "status": "error",
            "stderr": (
                f"Language '{language}' is not supported. "
                "Only 'python' and 'bash' are currently implemented."
            ),
        }

    @staticmethod
    def _run(cmd, shell: bool = False, timeout: int = 30) -> dict:
        try:
            result = subprocess.run(
                cmd, shell=shell, capture_output=True, text=True, timeout=timeout
            )
            return {
                "status": "success" if result.returncode == 0 else "error",
                "stdout": result.stdout,
                "stderr": result.stderr,
            }
        except subprocess.TimeoutExpired:
            return {"status": "error", "stderr": f"Execution timed out after {timeout}s"}
        except Exception as e:
            return {"status": "error", "stderr": str(e)}


interpreter_instance = CodeInterpreter()


### =============== SEARCH TOOLS =============== ###


def _wiki_page_to_text(page, char_limit: int) -> str:
    """Convert a wikipedia.WikipediaPage to plain text, preserving table content."""
    from bs4 import BeautifulSoup
    soup = BeautifulSoup(page.html(), "html.parser")
    for tag in soup(["script", "style", "sup"]):
        tag.decompose()
    for table in soup.find_all("table"):
        rows = []
        for row in table.find_all("tr"):
            cells = [c.get_text(strip=True) for c in row.find_all(["th", "td"])]
            if any(cells):
                rows.append(" | ".join(cells))
        table.replace_with("\n".join(rows) + "\n")
    text = soup.get_text(separator="\n", strip=True)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text[:char_limit]


@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 results.
    Full article text is returned including discography tables, track listings, and other structured data.
    Args:
        query: The search query."""
    import wikipedia

    char_limit = 20000
    try:
        titles = wikipedia.search(query, results=2)
        if not titles:
            return "No Wikipedia results found."
        parts = []
        for title in titles:
            try:
                page = wikipedia.page(title, auto_suggest=False)
                text = _wiki_page_to_text(page, char_limit)
                truncated = len(text) >= char_limit
                notice = (
                    "\n[TRUNCATED: This article exceeds the retrieval limit. "
                    "Some information may be missing.]"
                    if truncated else ""
                )
                parts.append(
                    f'<Document source="{page.url}"/>\n{text}{notice}\n</Document>'
                )
            except wikipedia.exceptions.DisambiguationError as e:
                try:
                    page = wikipedia.page(e.options[0], auto_suggest=False)
                    text = _wiki_page_to_text(page, char_limit)
                    truncated = len(text) >= char_limit
                    notice = (
                        "\n[TRUNCATED: This article exceeds the retrieval limit. "
                        "Some information may be missing.]"
                        if truncated else ""
                    )
                    parts.append(
                        f'<Document source="{page.url}"/>\n{text}{notice}\n</Document>'
                    )
                except Exception:
                    pass
            except Exception as e:
                parts.append(
                    f'<Document source="https://en.wikipedia.org/wiki/{title}"/>\n'
                    f'Error loading page: {e}\n</Document>'
                )
        return "\n\n---\n\n".join(parts) if parts else "No Wikipedia results found."
    except Exception as e:
        return f"Error searching Wikipedia: {e}"


@tool
def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 3 results.
    Args:
        query: The search query."""
    try:
        response = TavilySearch(max_results=3).invoke(query)
        if isinstance(response, str):
            return response
        if isinstance(response, dict):
            if "error" in response:
                return f"Error searching Tavily: {response['error']}"
            results = response.get("results", [])
        else:
            results = response  # already a list of dicts (older versions)
        if not results:
            return "web_search returned no results."
        return "\n\n---\n\n".join(
            f'<Document source="{doc.get("url", "")}" title="{doc.get("title", "")}"/>\n{doc.get("content", "")}\n</Document>'
            for doc in results
        )
    except Exception as e:
        return f"Error searching Tavily: {e}"


@tool
def arxiv_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 results.
    Args:
        query: The search query."""
    try:
        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
        return "\n\n---\n\n".join(
            f'<Document source="{doc.metadata.get("entry_id", doc.metadata.get("source", ""))}" title="{doc.metadata.get("Title", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        )
    except Exception as e:
        return f"Error searching Arxiv: {e}"


### =============== CODE INTERPRETER TOOLS =============== ###


@tool
def execute_code_multilang(code: str, language: str = "python") -> str:
    """Execute code in Python or Bash and return the results.
    Args:
        code (str): The source code to execute.
        language (str): The language of the code. Supported: "python", "bash".
    Returns:
        A string summarising stdout, stderr, and execution status.
    """
    supported_languages = ["python", "bash"]
    language = language.lower()

    if language not in supported_languages:
        return f"Unsupported language: {language}. Supported: {', '.join(supported_languages)}"

    result = interpreter_instance.execute_code(code, language=language)
    response = []

    if result["status"] == "success":
        response.append(f"Code executed successfully in {language.upper()}")
        if result.get("stdout"):
            response.append("\nStandard Output:\n```\n" + result["stdout"].strip() + "\n```")
        if result.get("stderr"):
            response.append("\nStandard Error:\n```\n" + result["stderr"].strip() + "\n```")
    else:
        response.append(f"Code execution failed in {language.upper()}")
        if result.get("stderr"):
            response.append("\nError:\n```\n" + result["stderr"].strip() + "\n```")

    return "\n".join(response)


### =============== MATHEMATICAL TOOLS =============== ###


@tool
def multiply(a: float, b: float) -> float:
    """
    Multiplies two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a * b


@tool
def add(a: float, b: float) -> float:
    """
    Adds two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a + b


@tool
def subtract(a: float, b: float) -> float:
    """
    Subtracts two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a - b


@tool
def divide(a: float, b: float) -> float:
    """
    Divides two numbers.
    Args:
        a (float): the first float number
        b (float): the second float number
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b


@tool
def modulus(a: int, b: int) -> int:
    """
    Get the modulus of two numbers.
    Args:
        a (int): the first number
        b (int): the second number
    """
    return a % b


@tool
def power(a: float, b: float) -> float:
    """
    Get the power of two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a**b


@tool
def square_root(a: float) -> str:
    """
    Get the square root of a number. Returns a string to handle complex results.
    Args:
        a (float): the number to get the square root of
    """
    if a >= 0:
        return str(a**0.5)
    return str(cmath.sqrt(a))


### =============== DOCUMENT PROCESSING TOOLS =============== ###


@tool
def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
    """
    Save content to a file and return the path.
    Args:
        content (str): the content to save to the file
        filename (str, optional): the name of the file. If not provided, a random name file will be created.
    """
    temp_dir = tempfile.gettempdir()
    if filename is None:
        temp_file = tempfile.NamedTemporaryFile(delete=False, dir=temp_dir)
        filepath = temp_file.name
    else:
        filepath = os.path.join(temp_dir, filename)

    with open(filepath, "w") as f:
        f.write(content)

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


@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.
    Args:
        url (str): the URL of the file to download.
        filename (str, optional): the name of the file. If not provided, a random name file will be created.
    """
    try:
        if not filename:
            path = urlparse(url).path
            filename = os.path.basename(path)
            if not filename:
                filename = f"downloaded_{uuid.uuid4().hex[:8]}"

        temp_dir = tempfile.gettempdir()
        filepath = os.path.join(temp_dir, filename)

        response = requests.get(url, stream=True)
        response.raise_for_status()

        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 read this file to process its contents."
    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 OCR (pytesseract).
    Args:
        image_path (str): the path to the image file.
    """
    try:
        import pytesseract

        image = Image.open(image_path)
        text = pytesseract.image_to_string(image)
        return f"Extracted text from image:\n\n{text}"
    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.
    Args:
        file_path (str): the path to the CSV file.
        query (str): Question about the data
    """
    try:
        df = pd.read_csv(file_path)
        result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        result += f"Columns: {', '.join(df.columns)}\n\n"
        result += "Summary statistics:\n"
        result += str(df.describe())
        return result
    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.
    Args:
        file_path (str): the path to the Excel file.
        query (str): Question about the data
    """
    try:
        df = pd.read_excel(file_path)
        result = f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        result += f"Columns: {', '.join(df.columns)}\n\n"
        result += "Summary statistics:\n"
        result += str(df.describe())
        return result
    except Exception as e:
        return f"Error analyzing Excel file: {str(e)}"


### ============== IMAGE PROCESSING TOOLS =============== ###


@tool
def analyze_image(image_base64: str) -> Dict[str, Any]:
    """
    Analyze basic properties of an image (size, mode, color analysis, thumbnail preview).
    Args:
        image_base64 (str): Base64 encoded image string
    Returns:
        Dictionary with analysis result
    """
    try:
        img = decode_image(image_base64)
        width, height = img.size
        mode = img.mode

        if mode in ("RGB", "RGBA"):
            arr = np.array(img)
            avg_colors = arr.mean(axis=(0, 1))
            dominant = ["Red", "Green", "Blue"][np.argmax(avg_colors[:3])]
            brightness = avg_colors.mean()
            color_analysis = {
                "average_rgb": avg_colors.tolist(),
                "brightness": float(brightness),
                "dominant_color": dominant,
            }
        else:
            color_analysis = {"note": f"No color analysis for mode {mode}"}

        thumbnail = img.copy()
        thumbnail.thumbnail((100, 100))
        thumb_path = save_image(thumbnail, "thumbnails")
        thumbnail_base64 = encode_image(thumb_path)

        return {
            "dimensions": (width, height),
            "mode": mode,
            "color_analysis": color_analysis,
            "thumbnail": thumbnail_base64,
        }
    except Exception as e:
        return {"error": str(e)}


@tool
def transform_image(
    image_base64: str, operation: str, params: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Apply transformations: resize, rotate, crop, flip, brightness, contrast, blur, sharpen, grayscale.
    Args:
        image_base64 (str): Base64 encoded input image
        operation (str): Transformation operation
        params (Dict[str, Any], optional): Parameters for the operation
    Returns:
        Dictionary with transformed image (base64)
    """
    try:
        img = decode_image(image_base64)
        params = params or {}

        if operation == "resize":
            img = img.resize(
                (params.get("width", img.width // 2), params.get("height", img.height // 2))
            )
        elif operation == "rotate":
            img = img.rotate(params.get("angle", 90), expand=True)
        elif operation == "crop":
            img = img.crop(
                (
                    params.get("left", 0),
                    params.get("top", 0),
                    params.get("right", img.width),
                    params.get("bottom", img.height),
                )
            )
        elif operation == "flip":
            if params.get("direction", "horizontal") == "horizontal":
                img = img.transpose(Image.FLIP_LEFT_RIGHT)
            else:
                img = img.transpose(Image.FLIP_TOP_BOTTOM)
        elif operation == "adjust_brightness":
            img = ImageEnhance.Brightness(img).enhance(params.get("factor", 1.5))
        elif operation == "adjust_contrast":
            img = ImageEnhance.Contrast(img).enhance(params.get("factor", 1.5))
        elif operation == "blur":
            img = img.filter(ImageFilter.GaussianBlur(params.get("radius", 2)))
        elif operation == "sharpen":
            img = img.filter(ImageFilter.SHARPEN)
        elif operation == "grayscale":
            img = img.convert("L")
        else:
            return {"error": f"Unknown operation: {operation}"}

        result_path = save_image(img)
        return {"transformed_image": encode_image(result_path)}
    except Exception as e:
        return {"error": str(e)}


@tool
def draw_on_image(
    image_base64: str, drawing_type: str, params: Dict[str, Any]
) -> Dict[str, Any]:
    """
    Draw shapes (rectangle, circle, line) or text onto an image.
    Args:
        image_base64 (str): Base64 encoded input image
        drawing_type (str): Drawing type
        params (Dict[str, Any]): Drawing parameters
    Returns:
        Dictionary with result image (base64)
    """
    try:
        img = decode_image(image_base64)
        draw = ImageDraw.Draw(img)
        color = params.get("color", "red")

        if drawing_type == "rectangle":
            draw.rectangle(
                [params["left"], params["top"], params["right"], params["bottom"]],
                outline=color,
                width=params.get("width", 2),
            )
        elif drawing_type == "circle":
            x, y, r = params["x"], params["y"], params["radius"]
            draw.ellipse(
                (x - r, y - r, x + r, y + r), outline=color, width=params.get("width", 2)
            )
        elif drawing_type == "line":
            draw.line(
                (params["start_x"], params["start_y"], params["end_x"], params["end_y"]),
                fill=color,
                width=params.get("width", 2),
            )
        elif drawing_type == "text":
            font_size = params.get("font_size", 20)
            try:
                font = ImageFont.truetype("arial.ttf", font_size)
            except IOError:
                font = ImageFont.load_default()
            draw.text(
                (params["x"], params["y"]),
                params.get("text", "Text"),
                fill=color,
                font=font,
            )
        else:
            return {"error": f"Unknown drawing type: {drawing_type}"}

        result_path = save_image(img)
        return {"result_image": encode_image(result_path)}
    except Exception as e:
        return {"error": str(e)}


@tool
def generate_simple_image(
    image_type: str,
    width: int = 500,
    height: int = 500,
    params: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Generate a simple image (gradient, noise).
    Args:
        image_type (str): Type of image to generate ("gradient" or "noise")
        width (int): Image width in pixels
        height (int): Image height in pixels
        params (Dict[str, Any], optional): Specific parameters
    Returns:
        Dictionary with generated image (base64)
    """
    try:
        params = params or {}

        if image_type == "gradient":
            direction = params.get("direction", "horizontal")
            start_color = params.get("start_color", (255, 0, 0))
            end_color = params.get("end_color", (0, 0, 255))

            img = Image.new("RGB", (width, height))
            draw = ImageDraw.Draw(img)

            if direction == "horizontal":
                for x in range(width):
                    r = int(start_color[0] + (end_color[0] - start_color[0]) * x / width)
                    g = int(start_color[1] + (end_color[1] - start_color[1]) * x / width)
                    b = int(start_color[2] + (end_color[2] - start_color[2]) * x / width)
                    draw.line([(x, 0), (x, height)], fill=(r, g, b))
            else:
                for y in range(height):
                    r = int(start_color[0] + (end_color[0] - start_color[0]) * y / height)
                    g = int(start_color[1] + (end_color[1] - start_color[1]) * y / height)
                    b = int(start_color[2] + (end_color[2] - start_color[2]) * y / height)
                    draw.line([(0, y), (width, y)], fill=(r, g, b))

        elif image_type == "noise":
            noise_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
            img = Image.fromarray(noise_array, "RGB")

        else:
            return {"error": f"Unsupported image_type: {image_type}"}

        result_path = save_image(img)
        return {"generated_image": encode_image(result_path)}
    except Exception as e:
        return {"error": str(e)}


@tool
def combine_images(
    images_base64: List[str], operation: str, params: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Combine multiple images (stack horizontally or vertically).
    Args:
        images_base64 (List[str]): List of base64 images
        operation (str): Combination type ("stack")
        params (Dict[str, Any], optional): e.g. {"direction": "horizontal"} or "vertical"
    Returns:
        Dictionary with combined image (base64)
    """
    try:
        images = [decode_image(b64) for b64 in images_base64]
        params = params or {}

        if operation == "stack":
            direction = params.get("direction", "horizontal")
            if direction == "horizontal":
                total_width = sum(img.width for img in images)
                max_height = max(img.height for img in images)
                new_img = Image.new("RGB", (total_width, max_height))
                x = 0
                for img in images:
                    new_img.paste(img, (x, 0))
                    x += img.width
            else:
                max_width = max(img.width for img in images)
                total_height = sum(img.height for img in images)
                new_img = Image.new("RGB", (max_width, total_height))
                y = 0
                for img in images:
                    new_img.paste(img, (0, y))
                    y += img.height
        else:
            return {"error": f"Unsupported combination operation: {operation}"}

        result_path = save_image(new_img)
        return {"combined_image": encode_image(result_path)}
    except Exception as e:
        return {"error": str(e)}


def get_tools():
    """Return a list of available tools for the agent."""
    return [
        web_search,
        wiki_search,
        arxiv_search,
        multiply,
        add,
        subtract,
        divide,
        modulus,
        power,
        square_root,
        save_and_read_file,
        download_file_from_url,
        extract_text_from_image,
        analyze_csv_file,
        analyze_excel_file,
        execute_code_multilang,
        analyze_image,
        transform_image,
        draw_on_image,
        generate_simple_image,
        combine_images,
    ]