X-Agents/ai-core/parser/vlm_client.py

"""
VLM 客户端 - 用于调用 VLM 模型进行文档理解
"""
import logging
import base64
import requests
from typing import Optional, Dict, Any

logger = logging.getLogger(__name__)


class VLMClient:
    """VLM 客户端，支持多种提供商"""

    def __init__(self, config: Dict[str, Any]):
        """
        初始化 VLM 客户端

        Args:
            config: VLM 配置，包含 provider, model, api_key, base_url, prompt 等
        """
        self.config = config
        self.provider = config.get("provider", "openai")
        self.model = config.get("model", "gpt-4o")
        self.api_key = config.get("api_key", "")
        self.base_url = config.get("base_url", "")
        self.prompt = config.get("prompt", "") or self._default_prompt()

        logger.info(f"VLMClient initialized: provider={self.provider}, model={self.model}")

    def _default_prompt(self) -> str:
        """默认提示词"""
        return """请分析这张图片中的文档内容，并将其转换为 Markdown 格式。
要求：
1. 保持原文的格式和结构
2. 表格用 Markdown 表格格式
3. 标题用 # ## ### 标记
4. 代码块用 ``` 标记
5. 尽量保留原文的所有信息"""

    def analyze_image(self, image_data: bytes, mime_type: str = "image/png") -> Dict[str, Any]:
        """
        使用 VLM 分析图片

        Args:
            image_data: 图片二进制数据
            mime_type: 图片 MIME 类型

        Returns:
            包含分析结果的字典
        """
        if self.provider == "openai":
            return self._call_openai(image_data, mime_type)
        elif self.provider == "anthropic":
            return self._call_anthropic(image_data, mime_type)
        elif self.provider == "qwen":
            return self._call_qwen(image_data, mime_type)
        else:
            return {
                "success": False,
                "content": "",
                "error": f"Unsupported provider: {self.provider}"
            }

    def _call_openai(self, image_data: bytes, mime_type: str) -> Dict[str, Any]:
        """调用 OpenAI GPT-4o API"""
        try:
            url = (self.base_url or "https://api.openai.com/v1") + "/chat/completions"

            # Base64 编码图片
            image_base64 = base64.b64encode(image_data).decode("utf-8")
            data_url = f"data:{mime_type};base64,{image_base64}"

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

            payload = {
                "model": self.model,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": self.prompt},
                            {"type": "image_url", "image_url": {"url": data_url}}
                        ]
                    }
                ],
                "max_tokens": 4096
            }

            response = requests.post(url, headers=headers, json=payload, timeout=120)
            response.raise_for_status()

            result = response.json()
            content = result["choices"][0]["message"]["content"]

            return {
                "success": True,
                "content": content,
                "usage": result.get("usage", {})
            }
        except Exception as e:
            logger.error(f"OpenAI API error: {e}")
            return {
                "success": False,
                "content": "",
                "error": str(e)
            }

    def _call_anthropic(self, image_data: bytes, mime_type: str) -> Dict[str, Any]:
        """调用 Anthropic Claude API"""
        try:
            url = (self.base_url or "https://api.anthropic.com/v1") + "/messages"

            image_base64 = base64.b64encode(image_data).decode("utf-8")

            headers = {
                "x-api-key": self.api_key,
                "anthropic-version": "2023-06-01",
                "Content-Type": "application/json"
            }

            # Anthropic 支持 image 类型
            payload = {
                "model": self.model,
                "max_tokens": 4096,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": self.prompt},
                            {
                                "type": "image",
                                "source": {
                                    "type": "base64",
                                    "media_type": mime_type,
                                    "data": image_base64
                                }
                            }
                        ]
                    }
                ]
            }

            response = requests.post(url, headers=headers, json=payload, timeout=120)
            response.raise_for_status()

            result = response.json()
            content = result["content"][0]["text"]

            return {
                "success": True,
                "content": content,
                "usage": result.get("usage", {})
            }
        except Exception as e:
            logger.error(f"Anthropic API error: {e}")
            return {
                "success": False,
                "content": "",
                "error": str(e)
            }

    def _call_qwen(self, image_data: bytes, mime_type: str) -> Dict[str, Any]:
        """调用阿里 Qwen VL API"""
        try:
            url = (self.base_url or "https://dashscope.aliyuncs.com/compatible-mode/v1") + "/chat/completions"

            image_base64 = base64.b64encode(image_data).decode("utf-8")

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

            # Qwen 格式
            payload = {
                "model": self.model,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": self.prompt},
                            {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{image_base64}"}}
                        ]
                    }
                ]
            }

            response = requests.post(url, headers=headers, json=payload, timeout=120)
            response.raise_for_status()

            result = response.json()
            content = result["choices"][0]["message"]["content"]

            return {
                "success": True,
                "content": content,
                "usage": {}
            }
        except Exception as e:
            logger.error(f"Qwen API error: {e}")
            return {
                "success": False,
                "content": "",
                "error": str(e)
            }