201 lines
7.7 KiB
Python
201 lines
7.7 KiB
Python
"""可复用的 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()
|
||
}
|
||
from app.observability.tracing import observe_decision, record_decision
|
||
|
||
# 观测包在请求外面：成功才记结果，失败仍抛原来的 DecisionError
|
||
# generation 为 None 时 record_decision 什么都不做
|
||
with observe_decision(request) as generation:
|
||
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(generation, 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
|
||
):
|
||
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()
|