app/tools/image_tool.py

from __future__ import annotations

import base64
import re

import requests
from langchain_core.tools import tool

from app.services.vision.pipeline import (
    process_image_hybrid,
    process_image_force_vision,
)
from app.tools.scraper import _is_safe_url
from app.utils.logger import get_logger

logger = get_logger(__name__)

MAX_IMAGE_DOWNLOAD = 20 * 1024 * 1024


def _decode_base64(data_url: str) -> bytes:
    match = re.match(r"data:image/[^;]+;base64,(.+)", data_url)
    if not match:
        raise ValueError("Invalid base64 image data URL")
    return base64.b64decode(match.group(1))


def _fetch_image(url: str) -> bytes:
    if not _is_safe_url(url):
        raise ValueError("Image URL is not allowed. Only public HTTP/HTTPS URLs are supported.")
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
    }
    resp = requests.get(url, timeout=20, headers=headers, stream=True)
    resp.raise_for_status()
    content = b""
    for chunk in resp.iter_content(chunk_size=8192):
        content += chunk
        if len(content) > MAX_IMAGE_DOWNLOAD:
            raise ValueError(f"Image too large (max {MAX_IMAGE_DOWNLOAD // 1024 // 1024}MB)")
    return content


@tool("analyze_image")
def analyze_image(
    image_source: str,
    query: str = None,
    force_vision: str = "false",
) -> str:
    """Analyze an image using OCR and optional vision model for complex scenes."""
    try:
        if image_source.startswith("data:image"):
            image_bytes = _decode_base64(image_source)
            image_url = None
        elif image_source.startswith(("http://", "https://")):
            image_bytes = _fetch_image(image_source)
            image_url = image_source
        else:
            return "Error: image_source must be a valid URL or base64 data URL."

        force = force_vision.lower() == "true"
        
        if force:
            result = process_image_force_vision(
                image_bytes,
                image_url=image_url,
                custom_query=query,
            )
        else:
            result = process_image_hybrid(
                image_bytes,
                image_url=image_url,
                custom_query=query,
            )

        if result.complexity == "simple_text":
            return result.ocr_text if result.ocr_text else "No text found in image."
        else:
            output_parts = []
            if result.vision_analysis and "failed" not in result.vision_analysis.lower():
                output_parts.append(result.vision_analysis)
            if result.ocr_text and "cannot" not in result.ocr_text.lower():
                output_parts.append(f"Text in image: {result.ocr_text}")
            return "\n\n".join(output_parts) if output_parts else "Unable to analyze image."

    except requests.RequestException:
        return "Failed to fetch image. The URL may be unreachable or the image too large."
    except ValueError:
        return "Image analysis failed. The image format or URL is not supported."
    except Exception:
        logger.exception("Image tool error")
        return "Image analysis failed. Please try again."