hendrik/tools/vision.py

import os
import re
import io
import base64
import tempfile
import urllib.request
import ftplib
from urllib.parse import urlparse
from tools.coder import get_current_workspace

MIME_MAP = {
    ".jpg":  "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png":  "image/png",
    ".gif":  "image/gif",
    ".webp": "image/webp",
    ".bmp":  "image/bmp",
}

MIME_EXT_REVERSE = {v: k for k, v in MIME_MAP.items()}
MAX_SIZE = 10 * 1024 * 1024  # 10MB


class ImagePayload:
    def __init__(self, data, mime, path):
        self.data = data
        self.mime = mime
        self.path = path


schema_describe_image = {
    "type": "function",
    "function": {
        "name": "describe_image",
        "description": "Load an image and send it to the model for visual analysis. Supports local files, HTTP/HTTPS URLs, FTP, and SFTP paths.",
        "parameters": {
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": "Image source. Local path, or URL (http/https/ftp/sftp://user:pass@host/path)"
                }
            },
            "required": ["path"]
        }
    }
}


def _detect_mime(path):
    ext = os.path.splitext(path.split("?")[0])[1].lower()
    return MIME_MAP.get(ext)


def _validate_and_pack(raw, source_label):
    if len(raw) > MAX_SIZE:
        return f"Error: image too large ({len(raw) / 1024 / 1024:.1f}MB). Max: {MAX_SIZE / 1024 / 1024:.0f}MB"

    mime = _detect_mime(source_label)
    if not mime:
        return f"Error: cannot detect image format from '{source_label}'. Supported: {', '.join(MIME_MAP.keys())}"

    encoded = base64.b64encode(raw).decode("ascii")
    return ImagePayload(data=encoded, mime=mime, path=source_label)


def _fetch_http(url):
    req = urllib.request.Request(url, headers={"User-Agent": "Hendrik/1.0"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        raw = resp.read()
        # Try Content-Type header for MIME detection
        ct = resp.headers.get("Content-Type", "")
        ct_mime = ct.split(";")[0].strip().lower() if ct else None
        if ct_mime in MIME_EXT_REVERSE:
            # Append extension based on Content-Type for _detect_mime
            ext = MIME_EXT_REVERSE[ct_mime]
            final_url = resp.url + ext
        else:
            final_url = resp.url
        return raw, final_url


def _fetch_ftp(parsed):
    host = parsed.hostname
    port = parsed.port or 21
    path = parsed.path.lstrip("/")
    username = parsed.username or "anonymous"
    password = parsed.password or ""

    buf = io.BytesIO()
    ftp = ftplib.FTP()
    ftp.connect(host, port, timeout=30)
    ftp.login(username, password)
    ftp.retrbinary(f"RETR {path}", buf.write)
    ftp.quit()
    return buf.getvalue()


def _fetch_sftp(url, parsed):
    try:
        import paramiko
    except ImportError:
        return "Error: paramiko not installed. Install with: pip install paramiko"

    host = parsed.hostname
    port = parsed.port or 22
    path = parsed.path
    username = parsed.username or ""
    password = parsed.password or ""

    transport = paramiko.Transport((host, port))
    transport.connect(username=username, password=password)
    sftp = paramiko.SFTPClient.from_transport(transport)

    with sftp.open(path, "rb") as f:
        raw = f.read()

    sftp.close()
    transport.close()
    return raw


def describe_image(path):
    path = path.strip()

    # ── HTTP / HTTPS ───────────────────────────────────────────────────
    if re.match(r"https?://", path):
        try:
            raw, final_url = _fetch_http(path)
            return _validate_and_pack(raw, final_url)
        except Exception as e:
            return f"Error fetching URL: {e}"

    # ── FTP / FTPS ─────────────────────────────────────────────────────
    if re.match(r"ftps?://", path):
        parsed = urlparse(path)
        try:
            raw = _fetch_ftp(parsed)
            return _validate_and_pack(raw, path)
        except Exception as e:
            return f"Error fetching FTP: {e}"

    # ── SFTP ───────────────────────────────────────────────────────────
    if path.startswith("sftp://"):
        parsed = urlparse(path)
        try:
            raw = _fetch_sftp(path, parsed)
            if isinstance(raw, str):
                return raw  # error message
            return _validate_and_pack(raw, path)
        except Exception as e:
            return f"Error fetching SFTP: {e}"

    # ── Local file ─────────────────────────────────────────────────────
    full_path = os.path.join(get_current_workspace(), path) if not os.path.isabs(path) else path

    if not os.path.isfile(full_path):
        return f"Error: file not found: {full_path}"

    try:
        with open(full_path, "rb") as f:
            raw = f.read()
    except Exception as e:
        return f"Error reading file: {e}"

    return _validate_and_pack(raw, full_path)