ai-worker-platform/llm_gateway.py

# -*- coding: utf-8 -*-
"""
统一模型网关：多供应商 OpenAI 兼容协议调用 + 计量 + 计价
"""
import time
import requests
import config


class LLMError(Exception):
    pass


def get_provider_cfg(provider):
    cfg = config.PROVIDERS.get(provider)
    if not cfg:
        raise LLMError(f'未知供应商: {provider}')
    return cfg


def model_price(model):
    p = config.MODEL_PRICING.get(model, config.DEFAULT_PRICE)
    return p['input'], p['output']


def calc_cost(model, prompt_tokens, completion_tokens):
    pin, pout = model_price(model)
    return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6)


def chat(provider, model, messages, temperature=0.7, max_tokens=None,
         base_url=None, api_key=None, timeout=None, retries=None):
    """调用 OpenAI 兼容 chat/completions，返回 {text, usage, cost, model}
    messages 支持两种格式：
      - 纯文本：[{'role':'user','content':'...'}]
      - 多模态：[{'role':'user','content':[{'type':'text','text':'...'},
                                        {'type':'image_url','image_url':{'url':'...'}}]}]
    """
    cfg = get_provider_cfg(provider)
    url = (base_url or cfg['base_url']).rstrip('/') + '/chat/completions'
    key = api_key or cfg['api_key']
    if not key:
        raise LLMError(f'供应商 {provider} 未配置 API Key')
    headers = {
        'Authorization': f'Bearer {key}',
        'Content-Type': 'application/json',
    }
    payload = {
        'model': model,
        'messages': messages,
        'temperature': temperature,
    }
    if max_tokens:
        payload['max_tokens'] = max_tokens

    timeout = timeout or cfg.get('timeout', 300)
    retries = config.MAX_RETRY if retries is None else retries
    last_err = None
    for attempt in range(retries + 1):
        try:
            resp = requests.post(url, json=payload, headers=headers, timeout=timeout)
            if resp.status_code == 200:
                data = resp.json()
                msg = data['choices'][0]['message']
                text = msg.get('content') or ''
                if not text:
                    # 推理模型偶发 content 为空：用 reasoning_content 兜底
                    text = msg.get('reasoning_content') or ''
                if not text:
                    last_err = LLMError('模型返回空内容，重试中…')
                    continue
                usage = data.get('usage', {})
                pt = usage.get('prompt_tokens', 0)
                ct = usage.get('completion_tokens', 0)
                return {
                    'text': text,
                    'model': data.get('model', model),
                    'prompt_tokens': pt,
                    'completion_tokens': ct,
                    'total_tokens': pt + ct,
                    'cost': calc_cost(model, pt, ct),
                }
            if resp.status_code == 429:
                last_err = LLMError(f'模型限流(429): {resp.text[:200]}')
                time.sleep(2 * (attempt + 1))
                continue
            if resp.status_code >= 500:
                last_err = LLMError(f'服务端错误({resp.status_code}): {resp.text[:200]}')
                time.sleep(1)
                continue
            raise LLMError(f'调用失败({resp.status_code}): {resp.text[:300]}')
        except requests.exceptions.Timeout:
            last_err = LLMError(f'请求超时({timeout}s)')
        except requests.exceptions.ConnectionError as e:
            last_err = LLMError(f'连接失败: {e}')
    raise last_err or LLMError('未知错误')


def chat_vision(provider, model, text, image_url=None, image_path=None,
                temperature=0.4, max_tokens=2000, base_url=None, api_key=None,
                retries=3):
    """多模态视觉调用：文本 + 图片（URL 或本地路径/base64）。
    返回与 chat() 相同结构。
    容错：聚合 API 偶发路由到纯文本后端（不认识 image_url），自动重试。"""
    import base64 as _b64
    import time as _time
    content = [{'type': 'text', 'text': text}]
    img_url = image_url
    if image_path:
        with open(image_path, 'rb') as f:
            raw = f.read()
        mime = 'image/png'
        if image_path.lower().endswith(('.jpg', '.jpeg')):
            mime = 'image/jpeg'
        elif image_path.lower().endswith('.gif'):
            mime = 'image/gif'
        elif image_path.lower().endswith('.webp'):
            mime = 'image/webp'
        img_url = f'data:{mime};base64,{_b64.b64encode(raw).decode()}'
    if img_url:
        content.append({'type': 'image_url', 'image_url': {'url': img_url}})
    messages = [{'role': 'user', 'content': content}]
    last_err = None
    for attempt in range(max(1, retries)):
        try:
            return chat(provider, model, messages,
                        temperature=temperature, max_tokens=max_tokens,
                        base_url=base_url, api_key=api_key)
        except LLMError as e:
            last_err = e
            msg = str(e)
            # 仅对“多模态格式不被支持/图片无效”类错误重试（聚合后端路由问题）
            if any(k in msg for k in ('image_url', 'InvalidParameter', 'invalid_parameter',
                                      'does not appear to be valid', 'image')):
                _time.sleep(2 * (attempt + 1))
                continue
            raise
    raise last_err or LLMError('视觉调用失败')


def test_connection(provider, model, base_url=None, api_key=None):
    """连通性测试：发一条最小请求"""
    t0 = time.time()
    r = chat(provider, model,
             [{'role': 'user', 'content': '请回复"OK"两个字'}],
             temperature=0, max_tokens=16,
             base_url=base_url, api_key=api_key, timeout=30, retries=0)
    return {'ok': True, 'latency_ms': int((time.time() - t0) * 1000),
            'reply': r['text'][:50], 'cost': r['cost']}