Profeto/src/llm/provider.py

"""多提供商 LLM 抽象(OpenAI-compatible 接口)。

支持: OpenAI / Deepseek / Ollama / 任何 OpenAI-compatible 网关。
"""
from __future__ import annotations

import json
import logging
import time
from dataclasses import dataclass, field
from typing import Any

import httpx

from src.core.config import settings
from src.core.http_client import get_client
from src.core.runtime_config import get_runtime_value

logger = logging.getLogger(__name__)


@dataclass
class LLMResponse:
    content: str
    parsed: dict | None = None
    prompt_tokens: int | None = None
    completion_tokens: int | None = None
    latency_ms: int | None = None
    raw: dict | None = None
    error: str | None = None


@dataclass
class LLMProvider:
    """OpenAI-compatible async provider。"""

    api_key: str = ""
    base_url: str = "https://api.openai.com/v1"
    model: str = "gpt-4o"
    timeout: int = 60
    extra_headers: dict = field(default_factory=dict)

    async def chat(
        self,
        system: str,
        user: str,
        *,
        json_mode: bool = True,
        temperature: float = 0.3,
        max_tokens: int = 1000,
    ) -> LLMResponse:
        """发请求,返回结构化响应。"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            **self.extra_headers,
        }
        payload: dict[str, Any] = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if json_mode:
            payload["response_format"] = {"type": "json_object"}

        start = time.perf_counter()
        try:
            client = get_client()
            # 连接与读取分离:端点不可达时 10s 内快速失败,
            # 避免每个 agent 各挂满 LLM_TIMEOUT 导致整次预测长时间无响应
            resp = await client.post(
                f"{self.base_url}/chat/completions",
                headers=headers,
                json=payload,
                timeout=httpx.Timeout(connect=10.0, read=float(self.timeout), write=float(self.timeout), pool=10.0),
            )
            resp.raise_for_status()
            data = resp.json()
            latency = int((time.perf_counter() - start) * 1000)
            usage = data.get("usage", {})
            message = data["choices"][0]["message"]
            content = message.get("content") or ""
            if not content:
                # 推理模型可能把 token 全花在 reasoning_content 上
                raise RuntimeError(
                    "模型未返回文本内容"
                    + ("(token 花在推理上,请增大 max_tokens)" if message.get("reasoning_content") else "")
                )
            parsed = None
            parse_error: str | None = None
            if json_mode:
                try:
                    parsed = json.loads(content)
                except json.JSONDecodeError:
                    # 尝试从代码块提取
                    import re
                    m = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", content)
                    if m:
                        try:
                            parsed = json.loads(m.group(1))
                        except json.JSONDecodeError:
                            pass
                    if parsed is None:
                        # P0-3: JSON 解析失败必须显式报错,不能静默继续
                        parse_error = f"JSON parse failed: {content[:200]!r}"
                        logger.warning(parse_error)
            return LLMResponse(
                content=content,
                parsed=parsed,
                prompt_tokens=usage.get("prompt_tokens"),
                completion_tokens=usage.get("completion_tokens"),
                latency_ms=latency,
                raw=data,
                error=parse_error if parse_error else None,
            )
        except Exception as e:
            latency = int((time.perf_counter() - start) * 1000)
            logger.error("LLM request failed: %s", e)
            return LLMResponse(content="", error=str(e), latency_ms=latency)


async def get_default_provider() -> LLMProvider:
    """构造默认 provider:运行时配置(DB)优先,回落 .env。"""
    return LLMProvider(
        api_key=await get_runtime_value("LLM_API_KEY"),
        base_url=await get_runtime_value("LLM_BASE_URL"),
        model=await get_runtime_value("LLM_MODEL"),
        timeout=settings.LLM_TIMEOUT,
    )