xsl
c97f8ae3e1
CI / lint / typecheck / import-linter / tests / secrets (push) Canceled after 0s
ob_ai_call is the only input to V2 pricing, and it was dropping the two numbers needed to interpret its own cost column. tokens_reasoning was captured from the provider response but never written to the table, so the cost of thinking mode was invisible once it left the process. Folding it into tokens_out was not an option either -- it is billed at the output rate, and mixing the two hides exactly the difference we now have to watch. rate_tier records which price list applied. Without it a row's cost_cents cannot be rechecked: DeepSeek's peak rate is double off-peak from 2026-08-17, so recomputing an old row at "the current rate" silently produces a different answer than what was actually billed. Migration 0002 adds both columns. tokens_reasoning defaults to 0 rather than NULL so cost aggregates do not collapse to NULL on historical rows; rate_tier stays nullable because rows written before this change genuinely have an unknown tier, and NULL says that more honestly than a guessed 'flat'. The smoke script now asserts both: reasoning tokens must be 0 (a non-zero value means a provider default re-enabled thinking mode behind our backs) and rate_tier must be present. Verified against the real API -- tokens_reasoning=0, rate_tier=off_peak, 3.9s round trip.
361 lines
13 KiB
Python
Executable File
361 lines
13 KiB
Python
Executable File
"""The LLM gateway (架构文档 §8.1).
|
||
|
||
Built-in capabilities, all exercised on every call:
|
||
|
||
=================== =========================================================
|
||
场景化路由 ``scene → (provider, model, temperature, max_tokens, timeout)``
|
||
故障转移 + 熔断 primary 超时/5xx/限流 → 自动切备用；连续失败熔断，冷却后半开
|
||
结构化输出 传入 Pydantic Schema → JSON 约束 + 解析失败重试 + 校验
|
||
Prompt 版本 模板文件按 id+version 加载，版本号随每次调用落库
|
||
缓存前缀纪律 固定前缀含时间戳/随机 ID/昵称 → 拒绝或告警 (§7.2, TD-04)
|
||
成本计量 每次 provider 调用（含失败与重试）写一行 ``ob_ai_call``
|
||
预算保护 平台级日成本硬熔断：超限即拒绝非关键场景 (§13)
|
||
=================== =========================================================
|
||
|
||
Deliberately not here yet: streaming (SSE), 审核联动 (input/output moderation),
|
||
per-user rate limiting, A/B prompt split by user hash. See packages/llm/README.md.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import time
|
||
from collections.abc import Callable, Mapping, Sequence
|
||
from decimal import Decimal
|
||
|
||
from pydantic import BaseModel, ValidationError
|
||
|
||
from packages.domain.ids import next_id
|
||
from packages.llm.budget import DailyBudgetGuard
|
||
from packages.llm.circuit import CircuitBreakerRegistry
|
||
from packages.llm.cost import AICallRecord, AICallRecorder, RateResolver, compute_cost_cents
|
||
from packages.llm.errors import (
|
||
AllProvidersFailedError,
|
||
CircuitOpenError,
|
||
ProviderError,
|
||
SceneNotRoutedError,
|
||
StructuredOutputError,
|
||
)
|
||
from packages.llm.prompts import PromptRegistry, validate_cache_prefix
|
||
from packages.llm.providers.base import LLMProvider
|
||
from packages.llm.routing import SceneRouter
|
||
from packages.llm.types import (
|
||
Completion,
|
||
Message,
|
||
ProviderRequest,
|
||
Scene,
|
||
SceneConfig,
|
||
Usage,
|
||
)
|
||
from packages.shared.logging import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
_JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL)
|
||
_STRUCTURED_INVALID = "llm_structured_output_invalid"
|
||
|
||
|
||
def extract_json(text: str) -> str:
|
||
"""Pull the JSON object out of a model reply that may be fenced or chatty."""
|
||
match = _JSON_FENCE_RE.search(text)
|
||
candidate = (match.group(1) if match else text).strip()
|
||
start, end = candidate.find("{"), candidate.rfind("}")
|
||
if start != -1 and end > start:
|
||
return candidate[start : end + 1]
|
||
return candidate
|
||
|
||
|
||
class LLMGateway:
|
||
"""Single entry point for every LLM call in the platform."""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
providers: Mapping[str, LLMProvider],
|
||
router: SceneRouter,
|
||
prompts: PromptRegistry,
|
||
rates: RateResolver,
|
||
recorder: AICallRecorder,
|
||
budget: DailyBudgetGuard,
|
||
breakers: CircuitBreakerRegistry | None = None,
|
||
id_factory: Callable[[], int] = next_id,
|
||
timer: Callable[[], float] = time.perf_counter,
|
||
) -> None:
|
||
self._providers = dict(providers)
|
||
self._router = router
|
||
self._prompts = prompts
|
||
self._rates = rates
|
||
self._recorder = recorder
|
||
self._budget = budget
|
||
self._breakers = breakers or CircuitBreakerRegistry()
|
||
self._id_factory = id_factory
|
||
self._timer = timer
|
||
|
||
# -- public API ---------------------------------------------------------
|
||
|
||
async def complete(
|
||
self,
|
||
*,
|
||
scene: Scene,
|
||
messages: Sequence[Message] = (),
|
||
prompt_vars: Mapping[str, str] | None = None,
|
||
schema: type[BaseModel] | None = None,
|
||
user_id: int | None = None,
|
||
ref_type: str | None = None,
|
||
ref_id: str | None = None,
|
||
) -> Completion:
|
||
config = self._router.resolve(scene)
|
||
|
||
# Gate 1: platform-wide daily cost breaker, checked before any spend.
|
||
await self._budget.ensure_within_budget(scene=scene.value, critical=config.critical)
|
||
|
||
template = self._prompts.get(config.prompt_id, config.prompt_version)
|
||
violations = validate_cache_prefix(
|
||
template.system_prefix, mode=self._prompts.cache_prefix_mode, ref=template.ref
|
||
)
|
||
if violations:
|
||
logger.warning(
|
||
"llm_prompt_cache_prefix_violation",
|
||
prompt=template.ref,
|
||
violations=[str(v) for v in violations],
|
||
)
|
||
|
||
composed = self._compose(
|
||
system_prefix=template.system_prefix,
|
||
user_text=template.render_user(dict(prompt_vars or {})),
|
||
messages=messages,
|
||
schema=schema,
|
||
)
|
||
|
||
failures: list[Exception] = []
|
||
for provider_name in config.providers:
|
||
provider = self._providers.get(provider_name)
|
||
if provider is None:
|
||
failures.append(
|
||
SceneNotRoutedError(f"provider '{provider_name}' is not registered")
|
||
)
|
||
continue
|
||
|
||
breaker = self._breakers.get(provider_name)
|
||
if not breaker.allow():
|
||
logger.warning(
|
||
"llm_circuit_open", provider=provider_name, state=breaker.state.value
|
||
)
|
||
failures.append(CircuitOpenError(provider_name, f"circuit {breaker.state.value}"))
|
||
continue
|
||
|
||
try:
|
||
completion = await self._call_provider(
|
||
provider=provider,
|
||
config=config,
|
||
messages=composed,
|
||
schema=schema,
|
||
user_id=user_id,
|
||
prompt_id=template.prompt_id,
|
||
prompt_version=template.version,
|
||
ref_type=ref_type,
|
||
ref_id=ref_id,
|
||
)
|
||
except ProviderError as exc:
|
||
if not exc.retryable:
|
||
# A malformed request would fail on the backup too.
|
||
raise
|
||
breaker.record_failure()
|
||
failures.append(exc)
|
||
logger.warning(
|
||
"llm_provider_failed",
|
||
provider=provider_name,
|
||
scene=scene.value,
|
||
error_code=exc.code,
|
||
circuit=breaker.state.value,
|
||
)
|
||
continue
|
||
except StructuredOutputError as exc:
|
||
breaker.record_failure()
|
||
failures.append(exc)
|
||
continue
|
||
|
||
breaker.record_success()
|
||
return completion
|
||
|
||
raise AllProvidersFailedError(scene.value, failures)
|
||
|
||
def breaker_states(self) -> dict[str, str]:
|
||
return self._breakers.snapshot()
|
||
|
||
async def aclose(self) -> None:
|
||
for provider in self._providers.values():
|
||
await provider.aclose()
|
||
|
||
# -- internals ----------------------------------------------------------
|
||
|
||
@staticmethod
|
||
def _compose(
|
||
*,
|
||
system_prefix: str,
|
||
user_text: str,
|
||
messages: Sequence[Message],
|
||
schema: type[BaseModel] | None,
|
||
) -> list[Message]:
|
||
"""架构文档 §7.2 ③: 固定前缀放最前以命中缓存，变量部分一律排在后面。"""
|
||
composed: list[Message] = [Message(role="system", content=system_prefix)]
|
||
composed.extend(messages)
|
||
if user_text:
|
||
composed.append(Message(role="user", content=user_text))
|
||
if schema is not None:
|
||
composed.append(Message(role="user", content=_schema_instruction(schema)))
|
||
return composed
|
||
|
||
async def _call_provider(
|
||
self,
|
||
*,
|
||
provider: LLMProvider,
|
||
config: SceneConfig,
|
||
messages: Sequence[Message],
|
||
schema: type[BaseModel] | None,
|
||
user_id: int | None,
|
||
prompt_id: str,
|
||
prompt_version: str,
|
||
ref_type: str | None,
|
||
ref_id: str | None,
|
||
) -> Completion:
|
||
model = config.model_for(provider.name, provider.default_model)
|
||
conversation: list[Message] = list(messages)
|
||
max_attempts = config.structured_max_retries + 1 if schema is not None else 1
|
||
schema_name = schema.__name__ if schema is not None else "<none>"
|
||
last_detail = "no attempt was made"
|
||
|
||
for attempt in range(1, max_attempts + 1):
|
||
request = ProviderRequest(
|
||
model=model,
|
||
messages=tuple(conversation),
|
||
temperature=config.temperature,
|
||
max_tokens=config.max_tokens,
|
||
timeout_seconds=config.timeout_seconds,
|
||
json_mode=schema is not None,
|
||
thinking=config.thinking,
|
||
)
|
||
call_id = self._id_factory()
|
||
started = self._timer()
|
||
|
||
try:
|
||
response = await provider.complete(request)
|
||
except ProviderError as exc:
|
||
await self._write_record(
|
||
AICallRecord(
|
||
call_id=call_id,
|
||
scene=config.scene.value,
|
||
provider=provider.name,
|
||
model=model,
|
||
success=False,
|
||
usage=Usage(),
|
||
latency_ms=self._elapsed_ms(started),
|
||
cost_cents=Decimal("0"),
|
||
user_id=user_id,
|
||
prompt_id=prompt_id,
|
||
prompt_version=prompt_version,
|
||
rate_tier=self._rates.current_tier(),
|
||
error_code=exc.code,
|
||
error_message=str(exc)[:1000],
|
||
ref_type=ref_type,
|
||
ref_id=ref_id,
|
||
)
|
||
)
|
||
raise
|
||
|
||
latency_ms = self._elapsed_ms(started)
|
||
# 单价与档位在同一刻取，成本落库后不再回算 —— 跨时段重算会改掉历史成本。
|
||
rate_tier = self._rates.current_tier()
|
||
cost = compute_cost_cents(response.usage, self._rates.for_model(response.model))
|
||
|
||
parsed: BaseModel | None = None
|
||
success = True
|
||
error_code: str | None = None
|
||
if schema is not None:
|
||
try:
|
||
parsed = schema.model_validate_json(extract_json(response.text))
|
||
except (ValidationError, ValueError) as exc:
|
||
success = False
|
||
error_code = _STRUCTURED_INVALID
|
||
last_detail = str(exc)[:400]
|
||
|
||
await self._write_record(
|
||
AICallRecord(
|
||
call_id=call_id,
|
||
scene=config.scene.value,
|
||
provider=provider.name,
|
||
model=response.model,
|
||
success=success,
|
||
usage=response.usage,
|
||
latency_ms=latency_ms,
|
||
cost_cents=cost,
|
||
user_id=user_id,
|
||
prompt_id=prompt_id,
|
||
prompt_version=prompt_version,
|
||
rate_tier=rate_tier,
|
||
error_code=error_code,
|
||
error_message=last_detail if not success else None,
|
||
ref_type=ref_type,
|
||
ref_id=ref_id,
|
||
)
|
||
)
|
||
# Budget accrues on every attempt, including ones we then discard —
|
||
# a validation-retry loop costs real money and must be visible.
|
||
await self._budget.add_cost(cost)
|
||
|
||
if success:
|
||
return Completion(
|
||
text=response.text,
|
||
scene=config.scene,
|
||
provider=provider.name,
|
||
model=response.model,
|
||
prompt_id=prompt_id,
|
||
prompt_version=prompt_version,
|
||
usage=response.usage,
|
||
latency_ms=latency_ms,
|
||
cost_cents=cost,
|
||
call_id=call_id,
|
||
parsed=parsed,
|
||
attempts=attempt,
|
||
)
|
||
|
||
logger.warning(
|
||
"llm_structured_output_invalid",
|
||
provider=provider.name,
|
||
scene=config.scene.value,
|
||
schema=schema_name,
|
||
attempt=attempt,
|
||
)
|
||
conversation.append(Message(role="assistant", content=response.text))
|
||
conversation.append(Message(role="user", content=_repair_instruction(last_detail)))
|
||
|
||
raise StructuredOutputError(provider.name, schema_name, max_attempts, last_detail)
|
||
|
||
def _elapsed_ms(self, started: float) -> int:
|
||
return max(int((self._timer() - started) * 1000), 0)
|
||
|
||
async def _write_record(self, record: AICallRecord) -> None:
|
||
"""Cost attribution is mandatory but must never break the user's request."""
|
||
try:
|
||
await self._recorder.record(record.with_defaults())
|
||
except Exception:
|
||
logger.exception("ai_call_record_failed", call_id=record.call_id, scene=record.scene)
|
||
|
||
|
||
def _schema_instruction(schema: type[BaseModel]) -> str:
|
||
import json
|
||
|
||
return (
|
||
"请只输出一个 JSON 对象，不要包含解释文字或代码块标记。"
|
||
"该对象必须严格符合以下 JSON Schema：\n"
|
||
f"{json.dumps(schema.model_json_schema(), ensure_ascii=False)}"
|
||
)
|
||
|
||
|
||
def _repair_instruction(detail: str) -> str:
|
||
return (
|
||
"上一次输出不符合要求的 JSON Schema，校验错误如下：\n"
|
||
f"{detail}\n"
|
||
"请重新只输出修正后的 JSON 对象，不要包含任何其他文字。"
|
||
)
|