open-swe/agent/utils/multimodal.py

"""Utilities for building multimodal content blocks."""

from __future__ import annotations

import base64
import logging
import mimetypes
import os
import re
from typing import Any
from urllib.parse import urlparse

import httpx
from langchain_core.messages.content import create_image_block

from .url_safety import request_with_safe_redirects

logger = logging.getLogger(__name__)

IMAGE_MARKDOWN_RE = re.compile(r"!\[[^\]]*\]\((https?://[^\s)]+)\)")
IMAGE_URL_RE = re.compile(
    r"(https?://[^\s)]+\.(?:png|jpe?g|gif|webp|bmp|tiff)(?:\?[^\s)]+)?)",
    re.IGNORECASE,
)


def extract_image_urls(text: str) -> list[str]:
    """Extract image URLs from markdown image syntax and direct image links."""
    if not text:
        return []

    urls: list[str] = []
    urls.extend(IMAGE_MARKDOWN_RE.findall(text))
    urls.extend(IMAGE_URL_RE.findall(text))

    deduped = dedupe_urls(urls)
    if deduped:
        logger.debug("Extracted %d image URL(s)", len(deduped))
    return deduped


def vision_not_supported_warning(model_id: str, image_count: int) -> str:
    """Build a prompt-visible warning when images are sent to a text-only model."""
    return (
        f"\n\n**Note:** {image_count} image(s) were attached but the current model "
        f"({model_id}) does not support image input. The images were not included. "
        "Please switch to a vision-enabled model to process images."
    )


async def fetch_image_block(
    image_url: str,
    client: httpx.AsyncClient,
) -> dict[str, Any] | None:
    """Fetch image bytes and build an image content block.

    The fetch validates and pins every redirect hop (SSRF guard) and drops any
    Authorization header once the URL redirects, so a per-host token is never
    replayed to a redirect target the caller never chose to authenticate to.
    URLs are logged host-only — a signed image URL can carry a bearer token.
    """
    host = (urlparse(image_url).hostname or "").lower()
    try:
        headers = None
        if host == "uploads.linear.app" or host.endswith(".uploads.linear.app"):
            linear_api_key = os.environ.get("LINEAR_API_KEY", "")
            if linear_api_key:
                headers = {"Authorization": linear_api_key}
            else:
                logger.warning(
                    "LINEAR_API_KEY not set; cannot authenticate image fetch for %s", host
                )
        elif host == "files.slack.com" or host.endswith(".files.slack.com"):
            slack_bot_token = os.environ.get("SLACK_BOT_TOKEN", "")
            if slack_bot_token:
                headers = {"Authorization": f"Bearer {slack_bot_token}"}
            else:
                logger.warning(
                    "SLACK_BOT_TOKEN not set; cannot authenticate image fetch for %s", host
                )

        response, blocked = await request_with_safe_redirects(
            client, "GET", image_url, headers=headers, strip_auth_on_redirect=True
        )
        if blocked is not None:
            _, reason = blocked
            logger.warning("Refusing to fetch image (SSRF guard) from %s: %s", host, reason)
            return None
        response.raise_for_status()
        content_type = response.headers.get("Content-Type", "").split(";")[0].strip()
        if not content_type:
            guessed, _ = mimetypes.guess_type(image_url)
            if not guessed:
                logger.warning("Could not determine content type from %s; skipping image", host)
                return None
            content_type = guessed

        supported_types = {"image/jpeg", "image/png", "image/gif", "image/webp"}
        if content_type not in supported_types:
            logger.warning(
                "Unsupported content type '%s' from %s; skipping image", content_type, host
            )
            return None

        encoded = base64.b64encode(response.content).decode("ascii")
        logger.info(
            "Fetched image from %s (%s, %d bytes)", host, content_type, len(response.content)
        )
        return create_image_block(base64=encoded, mime_type=content_type)
    except Exception:
        logger.exception("Failed to fetch image from %s", host)
        return None


def dedupe_urls(urls: list[str]) -> list[str]:
    return list(dict.fromkeys(urls))