ForcePilot/backend/package/yuxi/channels/adapters/yuanbao/vision.py

from __future__ import annotations

import base64
from dataclasses import dataclass
from typing import Any

import aiohttp

from yuxi.utils.logging_config import logger


@dataclass
class VisionResult:
    description: str
    confidence: float = 0.0
    raw_response: dict[str, Any] | None = None


async def analyze_image(
    image_data: bytes,
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
    prompt: str | None = None,
    mime_type: str = "image/jpeg",
) -> VisionResult | None:
    try:
        image_b64 = base64.b64encode(image_data).decode("utf-8")

        payload = {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{image_b64}",
                                "detail": "auto",
                            },
                        },
                        {
                            "type": "text",
                            "text": prompt or "请描述这张图片的内容",
                        },
                    ],
                }
            ],
            "model": "yuanbao-vision",
            "max_tokens": 500,
        }

        headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}

        async with http_client.post(
            f"{api_base}/api/v1/vision/analyze",
            json=payload,
            headers=headers,
            timeout=aiohttp.ClientTimeout(total=30),
        ) as resp:
            if resp.status == 200:
                data = await resp.json()
                description = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                return VisionResult(
                    description=description,
                    confidence=data.get("confidence", 0.0),
                    raw_response=data,
                )
            else:
                logger.warning(f"[Yuanbao] Vision analysis returned {resp.status}")
                return None
    except Exception as e:
        logger.warning(f"[Yuanbao] Vision analysis failed: {e}")
        return None


async def download_and_analyze(
    image_url: str,
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
    prompt: str | None = None,
) -> VisionResult | None:
    try:
        from urllib.parse import urlparse

        headers = {}
        trusted_hosts = {
            "open-api.yuanbao.tencent.com",
            "api.yuanbao.tencent.com",
        }
        parsed = urlparse(image_url)
        if parsed.hostname in trusted_hosts:
            headers["Authorization"] = f"Bearer {token}"

        async with http_client.get(
            image_url,
            headers=headers,
            timeout=aiohttp.ClientTimeout(total=30),
        ) as resp:
            if resp.status != 200:
                logger.warning(f"[Yuanbao] Image download failed: HTTP {resp.status}")
                return None

            image_data = await resp.read()
            content_type = resp.headers.get("Content-Type", "image/jpeg")

            return await analyze_image(
                image_data=image_data,
                api_base=api_base,
                token=token,
                http_client=http_client,
                prompt=prompt,
                mime_type=content_type,
            )
    except Exception as e:
        logger.warning(f"[Yuanbao] Image download and analyze failed: {e}")
        return None


async def analyze_images(
    image_urls: list[str],
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
    prompt: str | None = None,
    concurrency: int = 3,
) -> list[VisionResult | None]:
    import asyncio

    semaphore = asyncio.Semaphore(concurrency)

    async def _analyze_one(url: str) -> VisionResult | None:
        async with semaphore:
            return await download_and_analyze(
                image_url=url,
                api_base=api_base,
                token=token,
                http_client=http_client,
                prompt=prompt,
            )

    tasks = [_analyze_one(url) for url in image_urls]
    return await asyncio.gather(*tasks)


async def analyze_image_batch(
    images: list[tuple[bytes, str]],
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
    prompt: str | None = None,
    concurrency: int = 3,
) -> list[VisionResult | None]:
    import asyncio

    semaphore = asyncio.Semaphore(concurrency)

    async def _analyze_one(image_data: bytes, mime_type: str) -> VisionResult | None:
        async with semaphore:
            return await analyze_image(
                image_data=image_data,
                api_base=api_base,
                token=token,
                http_client=http_client,
                prompt=prompt,
                mime_type=mime_type,
            )

    tasks = [_analyze_one(data, mime) for data, mime in images]
    return await asyncio.gather(*tasks)