my-v0-server/app/clients/decision.py

"""可复用的 Jev 决策客户端；重试和 HTTP 连接管理交给官方 SDK"""

from functools import cache

from pydantic import ValidationError
from typesafe_sdk import (
    AsyncTypeSafeClient,
    RetryPolicy,
    SystemOneResponse,
    TypeSafeAPIConnectionError,
    TypeSafeAPIError,
    TypeSafeAPIResponseValidationError,
    TypeSafeAPITimeoutError,
    TypeSafeError,
)

from app.config import Settings, settings
from app.exceptions import DecisionError
from app.schemas.decision import (
    ChoiceDecision,
    ChoiceResult,
    DecisionRequest,
    DecisionResponse,
    ScoreDecision,
    ScoreResult,
)


class DecisionEngine:
    def __init__(self, config: Settings | None = None):
        self._config = config if config is not None else settings
        self._client: AsyncTypeSafeClient | None = None

    def _get_client(self) -> AsyncTypeSafeClient:
        if self._client is None:
            config = self._config
            if not config.typesafe_api_key.strip():
                raise DecisionError("请配置 TYPESAFE_API_KEY", kind="configuration")
            try:
                self._client = AsyncTypeSafeClient(
                    api_key=config.typesafe_api_key,
                    model=config.typesafe_model,
                    base_url=config.typesafe_base_url,
                    timeout=config.typesafe_timeout,
                    retry=RetryPolicy(
                        max_retries=config.typesafe_max_retries,
                        timeout=config.typesafe_timeout,
                    ),
                )
            except (TypeSafeError, ValueError) as error:
                raise DecisionError(
                    "Jev 客户端配置无效", kind="configuration"
                ) from error
        return self._client

    async def run(self, request: DecisionRequest) -> DecisionResponse:
        client = self._get_client()
        questions = {
            name: question.model_dump(mode="json", exclude_unset=True)
            | {"type": question.type}
            for name, question in request.decisions.items()
        }
        from app.observability.tracing import observe_decision, record_decision

        # 观测包在请求外面：成功才记结果，失败仍抛原来的 DecisionError
        # observation 为 None 时 record_decision 什么都不做
        with observe_decision(request) as observation:
            try:
                response = await client.system_one(
                    state=request.state,
                    questions=questions,
                    model=request.model,
                    extra_body=request.extra_body or None,
                )
            except TypeSafeAPIResponseValidationError as error:
                raise DecisionError(
                    "Jev 返回的数据格式无效",
                    kind="response_validation",
                    status=error.status,
                    request_id=error.request_id,
                ) from error
            except TypeSafeAPIError as error:
                raise DecisionError(
                    f"Jev 请求失败（HTTP {error.status}）",
                    kind="api",
                    status=error.status,
                    request_id=error.request_id,
                ) from error
            except TypeSafeAPITimeoutError as error:
                raise DecisionError("Jev 请求超时", kind="timeout") from error
            except TypeSafeAPIConnectionError as error:
                raise DecisionError("无法连接 Jev 服务", kind="connection") from error
            except TypeSafeError as error:
                raise DecisionError("Jev SDK 调用失败", kind="sdk") from error
            parsed = self._parse_response(request, response)
            record_decision(observation, parsed)
            return parsed

    @staticmethod
    def _parse_response(
        request: DecisionRequest, response: SystemOneResponse
    ) -> DecisionResponse:
        request_id = response.raw_http_response.headers.get("x-typesafe-request-id")
        try:
            result = DecisionResponse.model_validate(
                {
                    "results": {
                        name: answer.model_dump(mode="json")
                        for name, answer in response.answers.items()
                    },
                    "model": response.model,
                    "usage": response.usage.model_dump(mode="json"),
                    "request_id": request_id,
                }
            )
            if result.results.keys() != request.decisions.keys():
                raise ValueError("返回的问题 ID 与请求不一致")
            for name, question in request.decisions.items():
                answer = result.results[name]
                if answer.type != question.type:
                    raise ValueError("返回的决策类型与请求不一致")
                if (
                    isinstance(question, ChoiceDecision)
                    and isinstance(answer, ChoiceResult)
                    and answer.choice not in question.criteria
                ):
                    raise ValueError("返回了候选集合之外的选项")
                if (
                    isinstance(question, ScoreDecision)
                    and isinstance(answer, ScoreResult)
                    and not 0 <= answer.score <= len(question.criteria) - 1
                ):
                    raise ValueError("评分超出等级范围")
            return result
        except (ValidationError, ValueError) as error:
            raise DecisionError(
                "Jev 返回结果与决策协议不匹配",
                kind="response_validation",
                request_id=request_id,
            ) from error

    async def aclose(self) -> None:
        if self._client is not None:
            await self._client.aclose()
            self._client = None


@cache
def _cached_engine() -> DecisionEngine:
    return DecisionEngine()


def get_decision_engine() -> DecisionEngine:
    """与 get_llm 类似的共享入口；适用于单个服务事件循环"""
    return _cached_engine()


async def close_decision_engine() -> None:
    """在应用关闭阶段调用；跨 asyncio.run 时使用独立实例"""
    if _cached_engine.cache_info().currsize:
        await _cached_engine().aclose()
        _cached_engine.cache_clear()