feat: 封装 jev

2026-09-22 19:13:10 +08:00
parent 4d46336bc2
commit 94faffca18
2 changed files with 281 additions and 0 deletions
@@ -0,0 +1,193 @@
+"""可复用的 Jev 决策客户端；重试和 HTTP 连接管理交给官方 SDK"""
+
+from functools import cache
+from math import isclose
+from typing import Self
+
+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,
+        *,
+        client: AsyncTypeSafeClient | None = None,
+    ):
+        self._config = config if config is not None else settings
+        self._client = client
+        self._owns_client = client is None
+        self._closed = False
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    def _get_client(self) -> AsyncTypeSafeClient:
+        if self._closed:
+            raise DecisionError("决策客户端已关闭", kind="closed")
+        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:
+        # 即便模型实例在创建后被修改，也必须重新经过请求校验
+        request = DecisionRequest.model_validate(request.model_dump(mode="json"))
+        client = self._get_client()
+        questions = {
+            name: question.model_dump(mode="json", exclude_unset=True)
+            | {"type": question.type}
+            for name, question in request.decisions.items()
+        }
+        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
+        return self._parse_response(request, response)
+
+    @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
+                ):
+                    if answer.choice not in question.criteria:
+                        raise ValueError("返回了候选集合之外的选项")
+                    if answer.probabilities.keys() != question.criteria.keys():
+                        raise ValueError("选项概率不完整")
+                if isinstance(question, ScoreDecision) and isinstance(
+                    answer, ScoreResult
+                ):
+                    levels = {str(index) for index in range(len(question.criteria))}
+                    if (
+                        answer.probabilities.keys() != levels
+                        or answer.legend.keys() != levels
+                    ):
+                        raise ValueError("评分等级不完整")
+                    if not 0 <= answer.score <= len(question.criteria) - 1:
+                        raise ValueError("评分超出等级范围")
+                if isinstance(answer, (ChoiceResult, ScoreResult)):
+                    if not isclose(sum(answer.probabilities.values()), 1, abs_tol=1e-4):
+                        raise ValueError("概率之和不为 1")
+            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 not self._closed:
+            if self._owns_client and self._client is not None:
+                await self._client.aclose()
+            self._closed = True
+
+    async def __aenter__(self) -> Self:
+        self._get_client()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+        await self.aclose()
+
+
+@cache
+def _cached_engine() -> DecisionEngine:
+    return DecisionEngine()
+
+
+def get_decision_engine() -> DecisionEngine:
+    """与 get_llm 类似的共享入口；适用于单个服务事件循环"""
+    engine = _cached_engine()
+    if engine.closed:
+        _cached_engine.cache_clear()
+        engine = _cached_engine()
+    return engine
+
+
+async def close_decision_engine() -> None:
+    """在应用关闭阶段调用；跨 asyncio.run 时使用独立实例"""
+    if _cached_engine.cache_info().currsize:
+        await _cached_engine().aclose()
+        _cached_engine.cache_clear()