external_demand/app/core/open_router_llm.py

"""OpenRouter LLM 调用工具。"""

from __future__ import annotations

from typing import Any

from openrouter import OpenRouter
from openrouter.errors.openroutererror import OpenRouterError as SdkOpenRouterError

from app.core.config import settings

_ALLOWED_ROLES = frozenset({"system", "user", "assistant"})


class OpenRouterCallError(Exception):
    """项目内 OpenRouter 调用失败。"""

    def __init__(self, message: str, *, cause: Exception | None = None):
        super().__init__(message)
        self.cause = cause


def _build_client() -> OpenRouter:
    kwargs: dict[str, Any] = {
        "api_key": settings.open_router_api_key.strip(),
        "timeout_ms": max(settings.open_router_timeout_seconds, 1) * 1000,
    }
    if settings.open_router_http_referer.strip():
        kwargs["http_referer"] = settings.open_router_http_referer.strip()
    if settings.open_router_app_title.strip():
        kwargs["x_open_router_title"] = settings.open_router_app_title.strip()
    if settings.open_router_base_url.strip():
        kwargs["server_url"] = settings.open_router_base_url.strip()
    return OpenRouter(**kwargs)


def _normalize_messages(messages: list[dict[str, str]]) -> list[dict[str, str]]:
    if not messages:
        raise OpenRouterCallError("messages must not be empty")

    normalized: list[dict[str, str]] = []
    for index, item in enumerate(messages):
        role = (item.get("role") or "").strip()
        content = item.get("content")
        if role not in _ALLOWED_ROLES:
            raise OpenRouterCallError(
                f"messages[{index}].role must be one of: system, user, assistant"
            )
        if content is None or not str(content).strip():
            raise OpenRouterCallError(f"messages[{index}].content must not be empty")
        normalized.append({"role": role, "content": str(content)})
    return normalized


def _normalize_chat_response(response: Any) -> dict[str, Any]:
    choices = getattr(response, "choices", None) or []
    first_choice = choices[0] if choices else None
    message = getattr(first_choice, "message", None) if first_choice else None
    content = getattr(message, "content", None) or ""

    usage = getattr(response, "usage", None)
    if hasattr(usage, "model_dump"):
        usage = usage.model_dump()
    elif usage is not None and not isinstance(usage, dict):
        usage = dict(usage) if hasattr(usage, "__iter__") else usage

    return {
        "id": getattr(response, "id", None),
        "model": getattr(response, "model", None),
        "content": content,
        "usage": usage,
        "finish_reason": getattr(first_choice, "finish_reason", None) if first_choice else None,
    }


def create_chat_completion(
    messages: list[dict[str, str]],
    *,
    model: str | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
) -> dict[str, Any]:
    """通过 OpenRouter SDK 调用对话补全，返回结构化结果。"""
    if not settings.open_router_api_key.strip():
        raise OpenRouterCallError("OPEN_ROUTER_API_KEY or OPENROUTER_API_KEY is not configured")

    normalized_messages = _normalize_messages(messages)
    model_name = (model or settings.open_router_default_model).strip()
    if not model_name:
        raise OpenRouterCallError(
            "model is required (set OPEN_ROUTER_DEFAULT_MODEL or pass model=...)"
        )

    request_kwargs: dict[str, Any] = {
        "model": model_name,
        "messages": normalized_messages,
    }
    temperature_value = (
        temperature if temperature is not None else settings.open_router_temperature
    )
    max_tokens_value = max_tokens if max_tokens is not None else settings.open_router_max_tokens
    if temperature_value is not None:
        request_kwargs["temperature"] = temperature_value
    if max_tokens_value is not None:
        request_kwargs["max_tokens"] = max_tokens_value

    try:
        with _build_client() as client:
            response = client.chat.send(**request_kwargs)
    except SdkOpenRouterError as exc:
        raise OpenRouterCallError(str(exc), cause=exc) from exc
    except Exception as exc:
        raise OpenRouterCallError(f"OpenRouter SDK error: {exc}", cause=exc) from exc

    return _normalize_chat_response(response)


def chat_text(
    user_prompt: str,
    *,
    system_prompt: str | None = None,
    model: str | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
) -> str:
    """单轮对话便捷方法，直接返回模型回复文本。"""
    messages: list[dict[str, str]] = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    messages.append({"role": "user", "content": user_prompt.strip()})

    result = create_chat_completion(
        messages,
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return str(result.get("content") or "")