297 lines
10 KiB
Python
297 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
import uuid
|
|
from collections.abc import Iterator
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Protocol
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from backend.database.connection import Database
|
|
from backend.errors import BusinessError
|
|
from backend.llm.provider import OpenAICompatibleClient, ProviderFailure
|
|
from backend.llm.repository import LLMRepository
|
|
|
|
SHANGHAI = ZoneInfo("Asia/Shanghai")
|
|
|
|
|
|
class PrincipalUser(Protocol):
|
|
id: int
|
|
|
|
|
|
class PrincipalAccess(Protocol):
|
|
user: PrincipalUser
|
|
|
|
|
|
class MembershipViewAccess(Protocol):
|
|
quota_exempt: bool
|
|
daily_limit: int
|
|
|
|
|
|
class MembershipAccess(Protocol):
|
|
def view_for(self, principal: PrincipalAccess) -> MembershipViewAccess: ...
|
|
|
|
def can_use_smart_features(self, principal: PrincipalAccess) -> bool: ...
|
|
|
|
|
|
class ModelRecordAccess(Protocol):
|
|
id: int
|
|
base_url: str
|
|
model_identifier: str
|
|
|
|
|
|
class ModelRuntimeAccess(Protocol):
|
|
primary: ModelRecordAccess | None
|
|
fallback: ModelRecordAccess | None
|
|
|
|
|
|
class ModelPoolAccess(Protocol):
|
|
def runtime_config(self) -> ModelRuntimeAccess: ...
|
|
|
|
def decrypt_api_key(self, record: ModelRecordAccess) -> str: ...
|
|
|
|
|
|
class LLMGatewayError(RuntimeError):
|
|
def __init__(self, code: str, message: str, *, partial: bool = False) -> None:
|
|
super().__init__(message)
|
|
self.code = code
|
|
self.partial = partial
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class LLMProfile:
|
|
model_id: int
|
|
role: str
|
|
base_url: str
|
|
model_identifier: str
|
|
api_key: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class LLMCall:
|
|
request_id: str
|
|
user_id: int
|
|
feature: str
|
|
prompt_version: str
|
|
business_id: str
|
|
started_at: datetime
|
|
usage_date: str
|
|
input_chars: int
|
|
quota_exempt: bool
|
|
profiles: tuple[LLMProfile, ...]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class LLMStreamEvent:
|
|
type: str
|
|
content: str = ""
|
|
request_id: str = ""
|
|
|
|
|
|
class LLMGateway:
|
|
def __init__(
|
|
self,
|
|
database: Database,
|
|
repository: LLMRepository,
|
|
memberships: MembershipAccess,
|
|
model_pool: ModelPoolAccess,
|
|
provider: OpenAICompatibleClient | None = None,
|
|
) -> None:
|
|
self._database = database
|
|
self._repository = repository
|
|
self._memberships = memberships
|
|
self._model_pool = model_pool
|
|
self._provider = provider or OpenAICompatibleClient()
|
|
|
|
def prepare(
|
|
self,
|
|
principal: PrincipalAccess,
|
|
*,
|
|
feature: str,
|
|
prompt_version: str,
|
|
business_id: str,
|
|
input_chars: int,
|
|
) -> LLMCall:
|
|
membership = self._memberships.view_for(principal)
|
|
if not self._memberships.can_use_smart_features(principal):
|
|
raise LLMGatewayError("membership_required", "该功能仅对会员开放。")
|
|
try:
|
|
runtime = self._model_pool.runtime_config()
|
|
except BusinessError as exc:
|
|
raise LLMGatewayError("not_configured", "智能解读服务尚未配置。") from exc
|
|
records = (("primary", runtime.primary), ("fallback", runtime.fallback))
|
|
profiles = tuple(
|
|
LLMProfile(
|
|
model_id=record.id,
|
|
role=role,
|
|
base_url=record.base_url,
|
|
model_identifier=record.model_identifier,
|
|
api_key=self._model_pool.decrypt_api_key(record),
|
|
)
|
|
for role, record in records
|
|
if record is not None
|
|
)
|
|
now = datetime.now(SHANGHAI)
|
|
request_id = uuid.uuid4().hex
|
|
with self._database.transaction() as connection:
|
|
used = self._repository.usage_today(
|
|
connection, principal.user.id, now.date().isoformat()
|
|
)
|
|
active = self._repository.active_today(
|
|
connection, principal.user.id, now.date().isoformat()
|
|
)
|
|
if not membership.quota_exempt and used + active >= membership.daily_limit:
|
|
raise LLMGatewayError(
|
|
"quota_exhausted",
|
|
f"今日智能分析额度已用完（{membership.daily_limit} 次）。",
|
|
)
|
|
self._repository.reserve(
|
|
connection,
|
|
request_id=request_id,
|
|
user_id=principal.user.id,
|
|
feature=feature,
|
|
business_id=business_id,
|
|
prompt_version=prompt_version,
|
|
started_at=now.isoformat(timespec="seconds"),
|
|
input_chars=max(0, input_chars),
|
|
)
|
|
return LLMCall(
|
|
request_id=request_id,
|
|
user_id=principal.user.id,
|
|
feature=feature,
|
|
prompt_version=prompt_version,
|
|
business_id=business_id,
|
|
started_at=now,
|
|
usage_date=now.date().isoformat(),
|
|
input_chars=max(0, input_chars),
|
|
quota_exempt=membership.quota_exempt,
|
|
profiles=profiles,
|
|
)
|
|
|
|
def stream(
|
|
self, call: LLMCall, messages: list[dict[str, str]]
|
|
) -> Iterator[LLMStreamEvent]:
|
|
last_error = "unavailable"
|
|
for profile in call.profiles:
|
|
attempt_started = time.perf_counter()
|
|
with self._database.transaction() as connection:
|
|
attempt_id = self._repository.start_attempt(
|
|
connection,
|
|
call.request_id,
|
|
profile.model_id,
|
|
profile.role,
|
|
_now(),
|
|
call.input_chars,
|
|
)
|
|
output = ""
|
|
try:
|
|
upstream = iter(self._provider.stream(profile, messages))
|
|
first = next(upstream)
|
|
except (ProviderFailure, StopIteration) as exc:
|
|
last_error = exc.code if isinstance(exc, ProviderFailure) else "empty_response"
|
|
self._finish_attempt(attempt_id, "failed", attempt_started, last_error, 0)
|
|
continue
|
|
except Exception as exc:
|
|
last_error = "unavailable"
|
|
self._finish_attempt(attempt_id, "failed", attempt_started, last_error, 0)
|
|
self._finish_request(call, "failed", last_error, 0)
|
|
raise LLMGatewayError(
|
|
"model_unavailable", "智能解读服务暂不可用，请稍后重试。"
|
|
) from exc
|
|
with self._database.transaction() as connection:
|
|
self._repository.set_request_status(connection, call.request_id, "streaming")
|
|
output = first
|
|
try:
|
|
yield LLMStreamEvent("delta", first, call.request_id)
|
|
for chunk in upstream:
|
|
output += chunk
|
|
yield LLMStreamEvent("delta", chunk, call.request_id)
|
|
except GeneratorExit:
|
|
self._finish_attempt(
|
|
attempt_id, "stopped", attempt_started, "cancelled", len(output)
|
|
)
|
|
self._finish_request(call, "stopped", "cancelled", len(output))
|
|
raise
|
|
except ProviderFailure as exc:
|
|
self._finish_attempt(
|
|
attempt_id, "failed", attempt_started, exc.code, len(output)
|
|
)
|
|
self._finish_request(call, "failed", exc.code, len(output))
|
|
raise LLMGatewayError(
|
|
"stream_interrupted", "智能解读连接中断，请稍后重试。", partial=True
|
|
) from exc
|
|
except Exception as exc:
|
|
self._finish_attempt(
|
|
attempt_id, "failed", attempt_started, "unavailable", len(output)
|
|
)
|
|
self._finish_request(call, "failed", "unavailable", len(output))
|
|
raise LLMGatewayError(
|
|
"stream_interrupted", "智能解读连接中断，请稍后重试。", partial=True
|
|
) from exc
|
|
self._finish_attempt(attempt_id, "success", attempt_started, "", len(output))
|
|
self._finish_request(call, "success", "", len(output))
|
|
yield LLMStreamEvent("done", request_id=call.request_id)
|
|
return
|
|
self._finish_request(call, "failed", last_error, 0)
|
|
code, message = _safe_error(last_error)
|
|
raise LLMGatewayError(code, message)
|
|
|
|
def _finish_attempt(
|
|
self,
|
|
attempt_id: int,
|
|
status: str,
|
|
started: float,
|
|
error_type: str,
|
|
output_chars: int,
|
|
) -> None:
|
|
with self._database.transaction() as connection:
|
|
self._repository.finish_attempt(
|
|
connection,
|
|
attempt_id,
|
|
status,
|
|
_now(),
|
|
round((time.perf_counter() - started) * 1000),
|
|
error_type,
|
|
output_chars,
|
|
)
|
|
|
|
def _finish_request(
|
|
self, call: LLMCall, status: str, error_type: str, output_chars: int
|
|
) -> None:
|
|
now = datetime.now(SHANGHAI)
|
|
with self._database.transaction() as connection:
|
|
current = connection.execute(
|
|
"SELECT status FROM llm_requests WHERE id = ?", (call.request_id,)
|
|
).fetchone()
|
|
if current is None or str(current["status"]) in {"success", "failed", "stopped"}:
|
|
return
|
|
self._repository.set_request_status(
|
|
connection,
|
|
call.request_id,
|
|
status,
|
|
completed_at=now.isoformat(timespec="seconds"),
|
|
duration_ms=round((now - call.started_at).total_seconds() * 1000),
|
|
error_type=error_type,
|
|
output_chars=output_chars,
|
|
)
|
|
if status == "success" and not call.quota_exempt:
|
|
self._repository.increment_usage(
|
|
connection, call.user_id, call.usage_date, now.isoformat(timespec="seconds")
|
|
)
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(SHANGHAI).isoformat(timespec="seconds")
|
|
|
|
|
|
def _safe_error(error_type: str) -> tuple[str, str]:
|
|
messages = {
|
|
"capacity": ("model_capacity", "智能解读服务当前繁忙，请稍后重试。"),
|
|
"rate_limited": ("model_rate_limited", "智能解读请求过于频繁，请稍后重试。"),
|
|
"authentication": ("model_authentication", "智能解读服务配置失效，请联系管理员。"),
|
|
"timeout": ("model_timeout", "智能解读等待超时，请稍后重试。"),
|
|
"network": ("model_network", "智能解读网络暂不可用，请稍后重试。"),
|
|
}
|
|
return messages.get(error_type, ("model_unavailable", "智能解读服务暂不可用，请稍后重试。"))
|