llm_gateway.py

# -*- coding: utf-8 -*-
"""
统一模型网关（V3.5）：多供应商 OpenAI 兼容协议调用 + 流式 + 精细计量 + 计价
=====================================================================
核心变化：
1. 所有模型输出按 token 流式接收（SSE），不再一次性等完整响应；
2. 超时语义改为「单 token 返回超时 / 首字延迟超时 / 整体兜底」，
   三个值均可在 设置 页面配置（settings 表，动态生效）：
     - token_timeout        相邻两个 token 数据块的最大间隔（默认 60s）
     - first_token_timeout  请求发出后首块数据的最长等待（默认 120s）
     - request_timeout      整体兜底上限（默认 600s）
3. 大模型接口库（llm_endpoints）：base_url / api_key / 模型列表 / 逐模型定价，
   计费方式支持「按 token 数」与「按调用次数」两种，Worker 直接选用接口库；
4. 精细计量：每次调用记录 prompt / completion / 缓存命中 cached / 调用次数 /
   首字延迟 / 总耗时，全部入库（cost_records / agent_steps / chat_messages）。
"""
import json
import time
import requests
import config
import db


class LLMError(Exception):
    """模型调用异常。partial_text：超时中断前已收到的部分输出。"""

    def __init__(self, msg, partial_text=''):
        super().__init__(msg)
        self.partial_text = partial_text or ''


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


# ---------------------------------------------------------------------------
# 大模型接口库
# ---------------------------------------------------------------------------
def get_endpoint(endpoint_id):
    if not endpoint_id:
        return None
    try:
        return db.q('SELECT * FROM llm_endpoints WHERE id=? AND status="enabled"',
                    (int(endpoint_id),), one=True)
    except Exception:
        return None


def _endpoint_models(ep):
    try:
        return json.loads(ep.get('models') or '[]') or []
    except Exception:
        return []


def _endpoint_pricing_map(ep):
    try:
        return json.loads(ep.get('pricing') or '{}') or {}
    except Exception:
        return {}


def worker_llm_cfg(worker):
    """解析 Worker 的大模型接口配置（V3.5）：
    优先取绑定的接口库 endpoint_id（统一鉴权/计价），worker 自带 base_url/api_key 可覆盖。
    返回 dict(provider, model, base_url, api_key, endpoint, in_price, out_price,
              price_per_call, billing, pricing_map)"""
    ep = get_endpoint(worker.get('endpoint_id')) if worker else None
    if ep:
        models = _endpoint_models(ep)
        return {
            'provider': ep.get('provider') or 'custom',
            'model': worker.get('model') or (models[0] if models else ''),
            'base_url': worker.get('base_url') or ep.get('base_url') or '',
            'api_key': worker.get('api_key') or ep.get('api_key') or '',
            'endpoint': ep,
            'in_price': float(ep.get('input_price') or 0),
            'out_price': float(ep.get('output_price') or 0),
            'price_per_call': float(ep.get('price_per_call') or 0),
            'billing': ep.get('billing') or 'token',
            'pricing_map': _endpoint_pricing_map(ep),
        }
    return {
        'provider': (worker or {}).get('provider', ''),
        'model': (worker or {}).get('model', ''),
        'base_url': (worker or {}).get('base_url') or '',
        'api_key': (worker or {}).get('api_key') or '',
        'endpoint': None,
        'in_price': None, 'out_price': None, 'price_per_call': None,
        'billing': 'token', 'pricing_map': {},
    }


# ---------------------------------------------------------------------------
# 计价：优先接口库（逐模型定价 / 按调用次数），否则 config.MODEL_PRICING
# ---------------------------------------------------------------------------
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 calc_cost_ex(model, prompt_tokens, completion_tokens, calls=1, cfg=None):
    """按 Worker/接口库配置计价。cfg = worker_llm_cfg() 结果。"""
    if cfg and cfg.get('endpoint'):
        ep = cfg['endpoint']
        if ep.get('billing') == 'call':
            return round(float(ep.get('price_per_call') or 0) * max(1, calls), 6)
        p = (cfg.get('pricing_map') or {}).get(model)
        if p:
            pin, pout = float(p.get('input') or 0), float(p.get('output') or 0)
            if pin or pout:
                return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6)
        pin, pout = cfg.get('in_price') or 0, cfg.get('out_price') or 0
        if pin or pout:
            return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6)
    pin, pout = model_price(model)
    return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6)


def worker_unit_price(worker):
    """自动路由用：估算 Worker 单次调用成本（按 token 计费 = 输入价+0.5*输出价；按次计费 = 单价）"""
    cfg = worker_llm_cfg(worker)
    if cfg.get('endpoint') and cfg.get('billing') == 'call':
        return float(cfg.get('price_per_call') or 0)
    p = (cfg.get('pricing_map') or {}).get(cfg['model'])
    if p:
        return float(p.get('input') or 0) + float(p.get('output') or 0) * 0.5
    if cfg.get('in_price') is not None:
        return float(cfg['in_price']) + float(cfg['out_price']) * 0.5
    pin, pout = model_price(cfg['model'])
    return pin + pout * 0.5


# ---------------------------------------------------------------------------
# Token 估算（流式响应未带 usage 时的兜底）
# ---------------------------------------------------------------------------
def _estimate_prompt_tokens(messages):
    n = 0
    for m in messages or []:
        c = m.get('content') if isinstance(m, dict) else ''
        if isinstance(c, str):
            n += max(1, int(len(c) * 0.6))
        elif isinstance(c, list):
            for part in c:
                if not isinstance(part, dict):
                    continue
                t = part.get('text') or ''
                n += max(1, int(len(t) * 0.6))
                if part.get('image_url') or part.get('image'):
                    n += 1000  # 图片按约 1000 token 估算
    return max(1, n)


def _estimate_completion_tokens(text):
    return max(1, int(len(text or '') * 0.6))


# ---------------------------------------------------------------------------
# SSE 流式核心（单次尝试，无重试；超时语义见模块说明）
# ---------------------------------------------------------------------------
def _sock_of(resp):
    """尽力获取底层 socket 以调整读超时（兼容不同 requests/urllib3 版本）"""
    try:
        raw = resp.raw
        fp = getattr(raw, '_fp', None) or getattr(raw, 'fp', None)
        fpp = getattr(fp, 'fp', None)
        for obj in (fpp, fp):
            if obj is None:
                continue
            s = getattr(obj, 'raw', None) or getattr(obj, '_sock', None)
            if s is not None:
                return s
    except Exception:
        pass
    return None


def _chat_stream_raw(provider, model, messages, temperature=0.7, max_tokens=None,
                     base_url=None, api_key=None, token_timeout=60,
                     first_token_timeout=120, extra_payload=None):
    """发起一次流式请求，产出事件：
    yield ('delta', piece) | ('usage', usage_dict) | ('finish', finish_reason)
    结束前若收到 usage 则正常给出；异常抛 LLMError（含 partial_text）。
    超时：首块数据等待 first_token_timeout；相邻数据块间隔 token_timeout；整体 request_timeout 兜底。"""
    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',
        'Accept': 'text/event-stream',
    }
    payload = {
        'model': model,
        'messages': messages,
        'temperature': temperature,
        'stream': True,
        'stream_options': {'include_usage': True},
    }
    if max_tokens:
        payload['max_tokens'] = max_tokens
    if extra_payload:
        payload.update(extra_payload)

    t0 = time.time()
    first_token_at = None
    resp = None
    try:
        resp = requests.post(url, json=payload, headers=headers, stream=True,
                             timeout=(min(30, first_token_timeout), first_token_timeout))
        if resp.status_code != 200:
            body = resp.text[:300]
            resp.close()
            if resp.status_code == 429:
                raise LLMError(f'模型限流(429): {body}')
            if resp.status_code >= 500:
                raise LLMError(f'服务端错误({resp.status_code}): {body}')
            raise LLMError(f'调用失败({resp.status_code}): {body}')
        # 首块之后，读超时降为「单 token 返回超时」（首次 read 保持 first_token_timeout）
        sock = _sock_of(resp)
        first_line_seen = False
        for raw_line in resp.iter_lines(decode_unicode=True):
            line = (raw_line or '').strip()
            if not first_line_seen:
                first_line_seen = True
                if sock is not None:
                    try:
                        sock.settimeout(token_timeout)
                    except Exception:
                        pass
            if not line or not line.startswith('data:'):
                continue
            data = line[5:].strip()
            if data == '[DONE]':
                break
            try:
                evt = json.loads(data)
            except Exception:
                continue
            if evt.get('usage'):
                yield ('usage', evt['usage'])
                continue
            choices = evt.get('choices') or []
            if not choices:
                continue
            ch = choices[0]
            delta = ch.get('delta') or {}
            piece = delta.get('content') or ''
            if not piece:
                piece = delta.get('reasoning_content') or ''
            if piece:
                if first_token_at is None:
                    first_token_at = time.time()
                yield ('delta', piece)
            if ch.get('finish_reason'):
                yield ('finish', ch.get('finish_reason'))
        return
    except requests.exceptions.ReadTimeout:
        elapsed = time.time() - t0
        if first_token_at is None:
            raise LLMError(f'首字延迟超时（>{first_token_timeout}s 无输出）')
        raise LLMError(f'Token 返回超时（>{token_timeout}s 无新数据，已输出 {elapsed:.0f}s）')
    except requests.exceptions.Timeout:
        raise LLMError(f'请求超时（{(time.time() - t0):.0f}s）')
    except requests.exceptions.ConnectionError as e:
        raise LLMError(f'连接失败: {e}')
    finally:
        if resp is not None:
            try:
                resp.close()
            except Exception:
                pass


def _usage_fields(usage):
    """从 usage 中解析精细计量字段"""
    usage = usage or {}
    pt = int(usage.get('prompt_tokens') or 0)
    ct = int(usage.get('completion_tokens') or 0)
    det = usage.get('prompt_tokens_details') or {}
    cached = int(det.get('cached_tokens') or 0)
    if not cached:
        cached = int(usage.get('prompt_cache_hit_tokens') or 0)
    return pt, ct, cached


# ---------------------------------------------------------------------------
# 流式调用（带重试，聚合结果）
# ---------------------------------------------------------------------------
def chat_stream(provider, model, messages, temperature=0.7, max_tokens=None,
                base_url=None, api_key=None, timeout=None, retries=None,
                token_timeout=None, first_token_timeout=None, on_chunk=None,
                extra_payload=None):
    """流式接收全部输出，返回聚合结果 dict：
    {text, model, prompt_tokens, completion_tokens, total_tokens, cached_tokens,
     cost, calls, first_token_ms, elapsed_ms, usage_estimated, finish_reason}
    - 超时按「单 token 返回 / 首字延迟」语义（设置中可配）
    - 已有部分输出时不再重试（避免重复内容），否则按 retries 重试
    - on_chunk(delta) 逐块回调（用于对话流式转发）"""
    if token_timeout is None or first_token_timeout is None:
        tk, fk, rk = db.get_llm_timeouts()
        token_timeout = token_timeout or tk
        first_token_timeout = first_token_timeout or fk
    request_timeout = timeout or max(first_token_timeout + 5, 60)
    retries = config.MAX_RETRY if retries is None else retries
    last_err = None
    for attempt in range(retries + 1):
        parts, usage = [], None
        finish_reason = None
        first_token_at = None
        t0 = time.time()
        try:
            for evt, val in _chat_stream_raw(
                    provider, model, messages, temperature=temperature,
                    max_tokens=max_tokens, base_url=base_url, api_key=api_key,
                    token_timeout=token_timeout, first_token_timeout=first_token_timeout,
                    extra_payload=extra_payload):
                if evt == 'delta':
                    if first_token_at is None:
                        first_token_at = time.time()
                    parts.append(val)
                    if on_chunk:
                        try:
                            on_chunk(val)
                        except Exception:
                            pass
                elif evt == 'usage':
                    usage = val
                elif evt == 'finish':
                    finish_reason = val
            text = ''.join(parts)
            if not text and not usage:
                last_err = LLMError('模型返回空内容，重试中…')
                continue
            pt, ct, cached = _usage_fields(usage)
            usage_estimated = usage is None
            if usage is None:
                pt, ct = _estimate_prompt_tokens(messages), _estimate_completion_tokens(text)
            elapsed_ms = int((time.time() - t0) * 1000)
            first_ms = int((first_token_at - t0) * 1000) if first_token_at else None
            return {
                'text': text,
                'model': model,
                'prompt_tokens': pt,
                'completion_tokens': ct,
                'total_tokens': pt + ct,
                'cached_tokens': cached,
                'cost': calc_cost(model, pt, ct),
                'calls': 1,
                'first_token_ms': first_ms,
                'elapsed_ms': elapsed_ms,
                'usage_estimated': usage_estimated,
                'finish_reason': finish_reason,
            }
        except LLMError as e:
            last_err = e
            if e.partial_text:
                raise
            continue
    raise last_err or LLMError('未知错误')


# ---------------------------------------------------------------------------
# 兼容接口（非 Worker 场景）
# ---------------------------------------------------------------------------
def chat(provider, model, messages, temperature=0.7, max_tokens=None,
         base_url=None, api_key=None, timeout=None, retries=None,
         token_timeout=None, first_token_timeout=None, on_chunk=None, cfg=None):
    """兼容旧接口：流式接收全部输出后返回聚合结果。
    cfg = worker_llm_cfg() 结果时按接口库计价/鉴权。"""
    if cfg:
        provider = cfg['provider']
        base_url = cfg['base_url'] or None
        api_key = cfg['api_key'] or None
        model = cfg['model']
    r = chat_stream(provider, model, messages, temperature=temperature,
                    max_tokens=max_tokens, base_url=base_url, api_key=api_key,
                    timeout=timeout, retries=retries, token_timeout=token_timeout,
                    first_token_timeout=first_token_timeout, on_chunk=on_chunk)
    if cfg:
        r['cost'] = calc_cost_ex(model, r['prompt_tokens'], r['completion_tokens'], 1, cfg)
    return r


def chat_worker(worker, messages, temperature=None, max_tokens=None, on_chunk=None):
    """按 Worker 配置（含接口库）调用模型，返回聚合结果（含接口库计价与精细计量）"""
    cfg = worker_llm_cfg(worker)
    if not cfg['model']:
        raise LLMError(f'Worker「{worker.get("name", "")}」未配置模型')
    r = chat_stream(cfg['provider'], cfg['model'], messages,
                    temperature=temperature if temperature is not None else worker.get('temperature', 0.7),
                    max_tokens=max_tokens or worker.get('max_tokens') or 2000,
                    base_url=cfg['base_url'] or None, api_key=cfg['api_key'] or None,
                    on_chunk=on_chunk)
    r['cost'] = calc_cost_ex(cfg['model'], r['prompt_tokens'], r['completion_tokens'], 1, cfg)
    r['worker_id'] = worker['id']
    r['provider'] = cfg['provider']
    r['model'] = cfg['model']
    return r


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, endpoint=None):
    """连通性测试：发一条最小请求（流式），返回延迟/回复/成本"""
    t0 = time.time()
    cfg = None
    if endpoint:
        cfg = {'provider': endpoint.get('provider') or 'custom', 'model': model,
               'base_url': base_url or endpoint.get('base_url') or '',
               'api_key': api_key or endpoint.get('api_key') or '',
               'endpoint': endpoint}
    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,
             token_timeout=30, first_token_timeout=30, cfg=cfg)
    return {'ok': True, 'latency_ms': int((time.time() - t0) * 1000),
            'first_token_ms': r.get('first_token_ms'), 'reply': r['text'][:50],
            'cost': r['cost'], 'tokens': r['total_tokens']}