api_clients.py

import base64
import httpx
import asyncio
import logging
from pathlib import Path
from typing import Optional

from config import config

logger = logging.getLogger(__name__)


def image_to_base64(image_path: str) -> str:
    """将图片转为 base64"""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def get_mime_type(image_path: str) -> str:
    """获取图片 MIME 类型"""
    suffix = Path(image_path).suffix.lower()
    mime_map = {
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png": "image/png",
        ".gif": "image/gif",
        ".webp": "image/webp",
        ".bmp": "image/bmp",
    }
    return mime_map.get(suffix, "image/jpeg")


class BaseClient:
    """AI API 客户端基类"""
    
    async def analyze_image(self, image_path: str, custom_prompt: Optional[str] = None) -> str:
        raise NotImplementedError


class OpenAIClient(BaseClient):
    """OpenAI / OpenAI 兼容 API 客户端"""
    
    def __init__(self, api_key: str, base_url: str, model: str):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.model = model
    
    async def analyze_image(self, image_path: str, custom_prompt: Optional[str] = None) -> str:
        b64 = image_to_base64(image_path)
        mime = get_mime_type(image_path)
        prompt = custom_prompt or config.SYSTEM_PROMPT
        
        payload = {
            "model": self.model,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime};base64,{b64}",
                                "detail": "high"
                            }
                        }
                    ]
                }
            ],
            "max_tokens": 1024,
            "temperature": 0.3,
        }
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                f"{self.base_url}/chat/completions",
                json=payload,
                headers=headers
            )
            resp.raise_for_status()
            data = resp.json()
            return data["choices"][0]["message"]["content"].strip()


class GeminiClient(BaseClient):
    """Google Gemini API 客户端"""
    
    def __init__(self, api_key: str, model: str):
        self.api_key = api_key
        self.model = model
    
    async def analyze_image(self, image_path: str, custom_prompt: Optional[str] = None) -> str:
        b64 = image_to_base64(image_path)
        mime = get_mime_type(image_path)
        prompt = custom_prompt or config.SYSTEM_PROMPT
        
        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model}:generateContent"
        
        payload = {
            "contents": [
                {
                    "parts": [
                        {"text": prompt},
                        {
                            "inline_data": {
                                "mime_type": mime,
                                "data": b64
                            }
                        }
                    ]
                }
            ],
            "generationConfig": {
                "temperature": 0.3,
                "maxOutputTokens": 1024,
            }
        }
        
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                url,
                json=payload,
                params={"key": self.api_key}
            )
            resp.raise_for_status()
            data = resp.json()
            return data["candidates"][0]["content"]["parts"][0]["text"].strip()


class ClaudeClient(BaseClient):
    """Anthropic Claude API 客户端"""
    
    def __init__(self, api_key: str, model: str):
        self.api_key = api_key
        self.model = model
    
    async def analyze_image(self, image_path: str, custom_prompt: Optional[str] = None) -> str:
        b64 = image_to_base64(image_path)
        mime = get_mime_type(image_path)
        prompt = custom_prompt or config.SYSTEM_PROMPT
        
        payload = {
            "model": self.model,
            "max_tokens": 1024,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": mime,
                                "data": b64
                            }
                        },
                        {"type": "text", "text": prompt}
                    ]
                }
            ]
        }
        
        headers = {
            "x-api-key": self.api_key,
            "anthropic-version": "2023-06-01",
            "Content-Type": "application/json",
        }
        
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                "https://api.anthropic.com/v1/messages",
                json=payload,
                headers=headers
            )
            resp.raise_for_status()
            data = resp.json()
            return data["content"][0]["text"].strip()


def get_ai_client(
    provider: Optional[str] = None,
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    model: Optional[str] = None,
) -> BaseClient:
    """工厂方法：获取 AI 客户端"""
    provider = (provider or config.AI_PROVIDER).lower()
    
    if provider == "openai":
        return OpenAIClient(
            api_key=api_key or config.OPENAI_API_KEY,
            base_url=base_url or config.OPENAI_BASE_URL,
            model=model or config.OPENAI_MODEL,
        )
    elif provider == "gemini":
        return GeminiClient(
            api_key=api_key or config.GEMINI_API_KEY,
            model=model or config.GEMINI_MODEL,
        )
    elif provider == "claude":
        return ClaudeClient(
            api_key=api_key or config.CLAUDE_API_KEY,
            model=model or config.CLAUDE_MODEL,
        )
    elif provider == "qwen":
        return OpenAIClient(
            api_key=api_key or config.QWEN_API_KEY,
            base_url=base_url or config.QWEN_BASE_URL,
            model=model or config.QWEN_MODEL,
        )
    elif provider == "custom":
        return OpenAIClient(
            api_key=api_key or config.CUSTOM_API_KEY,
            base_url=base_url or config.CUSTOM_BASE_URL,
            model=model or config.CUSTOM_MODEL,
        )
    else:
        raise ValueError(f"不支持的 AI 提供商: {provider}")