diff --git a/backend/database/migrations/__init__.py b/backend/database/migrations/__init__.py index 19e24bd..b06796d 100644 --- a/backend/database/migrations/__init__.py +++ b/backend/database/migrations/__init__.py @@ -1,7 +1,8 @@ from .m0001_adopt_legacy import MIGRATION as M0001_ADOPT_LEGACY from .m0002_job_runs import MIGRATION as M0002_JOB_RUNS +from .m0003_llm_audit import MIGRATION as M0003_LLM_AUDIT from .runner import Migration, MigrationError, MigrationRunner -MIGRATIONS = (M0001_ADOPT_LEGACY, M0002_JOB_RUNS) +MIGRATIONS = (M0001_ADOPT_LEGACY, M0002_JOB_RUNS, M0003_LLM_AUDIT) __all__ = ["MIGRATIONS", "Migration", "MigrationError", "MigrationRunner"] diff --git a/backend/database/migrations/m0003_llm_audit.py b/backend/database/migrations/m0003_llm_audit.py new file mode 100644 index 0000000..efa51f1 --- /dev/null +++ b/backend/database/migrations/m0003_llm_audit.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import sqlite3 + +from backend.database.migrations.runner import Migration + + +def extend_llm_audit(connection: sqlite3.Connection) -> None: + columns = { + str(row["name"]) + for row in connection.execute("PRAGMA table_info(llm_usage)") + } + additions = ( + ("role", "TEXT NOT NULL DEFAULT ''"), + ("prompt_version", "TEXT NOT NULL DEFAULT ''"), + ("error_code", "TEXT NOT NULL DEFAULT ''"), + ("input_tokens", "INTEGER NOT NULL DEFAULT 0"), + ("output_tokens", "INTEGER NOT NULL DEFAULT 0"), + ) + for name, declaration in additions: + if name not in columns: + connection.execute( + f"ALTER TABLE llm_usage ADD COLUMN {name} {declaration}" + ) + + +MIGRATION = Migration( + version="0003", + name="extend_llm_audit", + action=extend_llm_audit, + signature="llm-audit:v1:role,prompt-version,error-code,input-tokens,output-tokens", +) diff --git a/backend/llm/__init__.py b/backend/llm/__init__.py new file mode 100644 index 0000000..c787fef --- /dev/null +++ b/backend/llm/__init__.py @@ -0,0 +1,15 @@ +from .gateway import ( + LLMGateway, + LLMGatewayError, + LLMResult, + LLMStreamEvent, + ModelProfile, +) + +__all__ = [ + "LLMGateway", + "LLMGatewayError", + "LLMResult", + "LLMStreamEvent", + "ModelProfile", +] diff --git a/backend/llm/gateway.py b/backend/llm/gateway.py new file mode 100644 index 0000000..d01f0c7 --- /dev/null +++ b/backend/llm/gateway.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +import time +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, Generic, TypeVar + + +T = TypeVar("T") + + +class LLMGatewayError(ValueError): + """Stable application error that does not expose provider details.""" + + def __init__(self, message: str, code: str = "unavailable") -> None: + super().__init__(message) + self.code = code + + +@dataclass(frozen=True) +class ModelProfile: + role: str + api_key: str + base_url: str + model: str + + @property + def configured(self) -> bool: + return bool(self.api_key and self.base_url and self.model) + + +@dataclass(frozen=True) +class LLMResult(Generic[T]): + value: T + source: str + role: str + model: str + latency_ms: int + + +@dataclass(frozen=True) +class LLMStreamEvent(Generic[T]): + kind: str + value: T | None = None + source: str = "" + role: str = "" + model: str = "" + latency_ms: int = 0 + + +class LLMGateway: + """Single policy boundary for access, model fallback, and call auditing.""" + + def __init__( + self, + *, + database: Any, + user_id_supplier: Callable[[], int], + membership_supplier: Callable[[], dict[str, Any]], + settings_supplier: Callable[[], dict[str, Any]], + profile_supplier: Callable[[], dict[str, Any]], + ) -> None: + self.database = database + self.user_id_supplier = user_id_supplier + self.membership_supplier = membership_supplier + self.settings_supplier = settings_supplier + self.profile_supplier = profile_supplier + + def ensure_access(self, feature: str) -> tuple[str, tuple[ModelProfile, ...]]: + del feature # Reserved for future feature-specific policy. + membership = self.membership_supplier() + profile = self.profile_supplier() + source = str(profile.get("source") or "none") + profiles = self._model_profiles(profile) + if source == "none" or not profiles: + raise LLMGatewayError("智能功能尚未配置，请联系管理员。", "not_configured") + if source == "platform": + settings = self.settings_supplier() + limit = max(1, int(settings.get("member_daily_limit") or 50)) + if not membership.get("active"): + raise LLMGatewayError("开通会员后可使用智能功能。", "membership_required") + if self._usage_today(source) >= limit: + raise LLMGatewayError( + f"今日会员模型额度已用完（{limit} 次）。", "quota_exhausted" + ) + return source, profiles + + def call( + self, + feature: str, + prompt_version: str, + invoke: Callable[[ModelProfile], T], + error_types: tuple[type[BaseException], ...], + ) -> LLMResult[T]: + source, profiles = self.ensure_access(feature) + started = time.perf_counter() + last_error: BaseException | None = None + for profile in profiles: + try: + value = invoke(profile) + except error_types as exc: + last_error = exc + continue + latency_ms = round((time.perf_counter() - started) * 1000) + self.audit( + feature, source, profile, "success", latency_ms, prompt_version + ) + return LLMResult( + value=value, + source=source, + role=profile.role, + model=profile.model, + latency_ms=latency_ms, + ) + failed = profiles[-1] + latency_ms = round((time.perf_counter() - started) * 1000) + self.audit( + feature, + source, + failed, + "failed", + latency_ms, + prompt_version, + self._error_code(last_error), + ) + raise LLMGatewayError("智能解读服务暂不可用，请稍后重试。") from last_error + + @staticmethod + def probe(profile: dict[str, Any], invoke: Callable[[ModelProfile], T]) -> T: + """Route an explicit administrator connection test through the gateway boundary.""" + model = ModelProfile( + role="probe", + api_key=str(profile.get("api_key") or ""), + base_url=str(profile.get("base_url") or ""), + model=str(profile.get("model") or ""), + ) + return invoke(model) + + def stream( + self, + feature: str, + prompt_version: str, + invoke: Callable[[ModelProfile], Iterator[T]], + error_types: tuple[type[BaseException], ...], + ) -> Iterator[LLMStreamEvent[T]]: + source, profiles = self.ensure_access(feature) + started = time.perf_counter() + last_error: BaseException | None = None + for profile in profiles: + try: + upstream = iter(invoke(profile)) + first = next(upstream) + except (*error_types, StopIteration) as exc: + last_error = exc + continue + yield LLMStreamEvent(kind="delta", value=first) + try: + for chunk in upstream: + yield LLMStreamEvent(kind="delta", value=chunk) + except error_types as exc: + latency_ms = round((time.perf_counter() - started) * 1000) + self.audit( + feature, + source, + profile, + "failed", + latency_ms, + prompt_version, + self._error_code(exc), + ) + raise LLMGatewayError( + "智能解读连接中断，请稍后重试。" + ) from exc + latency_ms = round((time.perf_counter() - started) * 1000) + self.audit( + feature, source, profile, "success", latency_ms, prompt_version + ) + yield LLMStreamEvent( + kind="complete", + source=source, + role=profile.role, + model=profile.model, + latency_ms=latency_ms, + ) + return + failed = profiles[-1] + latency_ms = round((time.perf_counter() - started) * 1000) + self.audit( + feature, + source, + failed, + "failed", + latency_ms, + prompt_version, + self._error_code(last_error), + ) + raise LLMGatewayError("智能解读服务暂不可用，请稍后重试。") from last_error + + def audit( + self, + feature: str, + source: str, + profile: ModelProfile, + status: str, + latency_ms: int, + prompt_version: str, + error_code: str = "", + input_tokens: int = 0, + output_tokens: int = 0, + ) -> None: + self.database.record_llm_usage( + self.user_id_supplier(), + feature, + source, + profile.model, + status, + latency_ms, + role=profile.role, + prompt_version=prompt_version, + error_code=error_code, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + + def _usage_today(self, source: str) -> int: + now = datetime.now().astimezone() + start = now.replace( + hour=0, minute=0, second=0, microsecond=0 + ).astimezone(timezone.utc) + return self.database.count_llm_usage_since( + self.user_id_supplier(), source, start.isoformat(timespec="seconds") + ) + + @staticmethod + def _model_profiles(profile: dict[str, Any]) -> tuple[ModelProfile, ...]: + result = [] + for role in ("primary", "fallback"): + item = profile.get(role) or {} + candidate = ModelProfile( + role=role, + api_key=str(item.get("api_key") or ""), + base_url=str(item.get("base_url") or ""), + model=str(item.get("model") or ""), + ) + if candidate.configured: + result.append(candidate) + return tuple(result) + + @staticmethod + def _error_code(error: BaseException | None) -> str: + if error is None: + return "empty_response" + return type(error).__name__[:80] diff --git a/database.py b/database.py index 51fac97..411276e 100644 --- a/database.py +++ b/database.py @@ -874,16 +874,26 @@ class ReviewDatabase: model: str, status: str, latency_ms: int = 0, + *, + role: str = "", + prompt_version: str = "", + error_code: str = "", + input_tokens: int = 0, + output_tokens: int = 0, ) -> None: now = datetime.now(timezone.utc).isoformat(timespec="seconds") with self.connect() as connection: connection.execute( """ INSERT INTO llm_usage - (user_id, feature, source, model, status, latency_ms, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) + (user_id, feature, source, model, status, latency_ms, created_at, + role, prompt_version, error_code, input_tokens, output_tokens) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, - (user_id, feature, source, model, status, int(latency_ms), now), + ( + user_id, feature, source, model, status, int(latency_ms), now, + role, prompt_version, error_code, int(input_tokens), int(output_tokens), + ), ) def count_llm_usage_since(self, user_id: int, source: str, since: str) -> int: diff --git a/docs/governance/architecture-inventory.json b/docs/governance/architecture-inventory.json index f73bee7..b6851bd 100644 --- a/docs/governance/architecture-inventory.json +++ b/docs/governance/architecture-inventory.json @@ -262,16 +262,16 @@ "bytes": 364385, "lines": 15549 }, - { - "path": "server.py", - "bytes": 268447, - "lines": 5956 - }, { "path": "static/redesign-v2.css", "bytes": 264960, "lines": 8616 }, + { + "path": "server.py", + "bytes": 263744, + "lines": 5856 + }, { "path": "static/index.html", "bytes": 133569, @@ -279,8 +279,8 @@ }, { "path": "database.py", - "bytes": 121153, - "lines": 2829 + "bytes": 121546, + "lines": 2839 }, { "path": "screener.py", diff --git a/docs/governance/stage-13-llm-gateway.md b/docs/governance/stage-13-llm-gateway.md new file mode 100644 index 0000000..2c454cc --- /dev/null +++ b/docs/governance/stage-13-llm-gateway.md @@ -0,0 +1,32 @@ +# Stage 13: Unified LLM Gateway + +## Boundary + +All runtime model calls now enter through `backend/llm/gateway.py`. Feature agents retain +their deterministic context assembly, prompt content, and provider response parsing. + +The gateway owns: + +- membership and daily quota enforcement; +- primary and fallback model selection; +- fallback only before the first streamed delta; +- stable user-visible availability and interruption errors; +- one logical-call audit record with feature, model role, prompt version, latency, status, + normalized error code, and token fields reserved for providers that report usage. + +Administrator connection probes also cross the gateway boundary, but do not consume member +quota or create usage records. + +## Compatibility + +- Mentor and review-assistant stream payloads are unchanged. +- Heaven readings retain their primary/fallback notice and persistence behavior. +- Strategy compilation still falls back to the deterministic local compiler when both model + profiles are unavailable, while access and quota failures remain blocking. +- Provider credentials and raw provider failures remain outside browser responses. + +## Residual Risk + +The current OpenAI-compatible streaming providers do not consistently return token usage, so +the audit schema records zero until transport adapters expose trustworthy token counts. Request +cancellation remains bounded by the existing provider socket timeout. diff --git a/server.py b/server.py index 0838c84..37fe92e 100644 --- a/server.py +++ b/server.py @@ -19,6 +19,7 @@ from assistant_agent import ReviewAssistantError, stream_review_assistant from api_access import ROUTES from backend.bootstrap import build_application_container, load_runtime_settings from backend.http import correlation_id, normalize_error_payload +from backend.llm import LLMGateway, LLMGatewayError from chart_data_provider import ChartDataError from app_config import ( DATA_DIR, @@ -193,6 +194,13 @@ class DashboardService: self.realtime_aggregator = self.container.realtime_aggregator self.chart_data = self.container.chart_data self.jobs = self.container.jobs + self.llm_gateway = LLMGateway( + database=self.database, + user_id_supplier=lambda: self.current_user_id, + membership_supplier=self.membership, + settings_supplier=lambda: self._system_credentials, + profile_supplier=self._resolved_llm_profile, + ) self.screener.ensure_builtin_strategies() self._background_stop = threading.Event() self._background_thread = self.jobs.start_scheduler( @@ -482,7 +490,12 @@ class DashboardService: raise ValueError("模型角色不支持。") profile = self._validate_llm_profile(payload, current, required=True, label=label) try: - return test_llm_connection(**profile) + return self.llm_gateway.probe( + profile, + lambda model: test_llm_connection( + model.api_key, model.base_url, model.model + ), + ) except LLMCompilerError as exc: raise ValueError(str(exc)) from exc @@ -533,29 +546,6 @@ class DashboardService: start.isoformat(timespec="seconds"), ) - def ensure_llm_access(self, feature: str) -> str: - profile = self._resolved_llm_profile() - source = str(profile.get("source") or "none") - if source == "none" or not self._profile_configured(profile.get("primary") or {}): - raise ValueError("智能功能尚未配置，请联系管理员。") - if source == "platform": - limit = max(1, int(self._system_credentials.get("member_daily_limit") or 50)) - if self._platform_usage_today() >= limit: - raise ValueError(f"今日会员模型额度已用完（{limit} 次）。") - return source - - def record_llm_usage( - self, - feature: str, - source: str, - model: str, - status: str, - latency_ms: int = 0, - ) -> None: - self.database.record_llm_usage( - self.current_user_id, feature, source, model, status, latency_ms - ) - def system_status(self) -> dict[str, Any]: platform = self._platform_llm_profile() model_pool = [] @@ -707,7 +697,12 @@ class DashboardService: payload, current, required=True, label=label ) try: - return test_llm_connection(**profile) + return self.llm_gateway.probe( + profile, + lambda model: test_llm_connection( + model.api_key, model.base_url, model.model + ), + ) except LLMCompilerError as exc: raise ValueError(str(exc)) from exc @@ -1617,55 +1612,33 @@ class DashboardService: for item in self.assistant_messages()[-12:] if item.get("role") in {"user", "assistant"} ] - source = self.ensure_llm_access("assistant") - profiles = [ - ( - self.llm_primary_api_key, - self.llm_primary_base_url, - self.llm_primary_model, - ) - ] - if self.llm_fallback_configured: - profiles.append( - ( - self.llm_fallback_api_key, - self.llm_fallback_base_url, - self.llm_fallback_model, - ) - ) - def generate(): - started = time.perf_counter() - last_error: Exception | None = None - for api_key, base_url, model in profiles: - try: - upstream = iter( - stream_review_assistant( - context, question, history, api_key, base_url, model - ) + answer_parts: list[str] = [] + events = self.llm_gateway.stream( + "assistant", + "review-assistant-v1", + lambda profile: stream_review_assistant( + context, + question, + history, + profile.api_key, + profile.base_url, + profile.model, + ), + (ReviewAssistantError,), + ) + for event in events: + if event.kind == "delta": + chunk = str(event.value or "") + answer_parts.append(chunk) + yield chunk + elif event.kind == "complete": + self.database.save_assistant_exchange( + self.current_user_id, + question, + "".join(answer_parts).strip(), + trade_date, ) - first = next(upstream) - except (ReviewAssistantError, StopIteration) as exc: - last_error = exc - continue - answer_parts = [first] - yield first - try: - for chunk in upstream: - answer_parts.append(chunk) - yield chunk - except ReviewAssistantError as exc: - self.record_llm_usage("assistant", source, model, "failed") - raise ValueError("智能解读连接中断，请稍后重试。") from exc - answer = "".join(answer_parts).strip() - latency_ms = round((time.perf_counter() - started) * 1000) - self.database.save_assistant_exchange( - self.current_user_id, question, answer, trade_date - ) - self.record_llm_usage("assistant", source, model, "success", latency_ms) - return - self.record_llm_usage("assistant", source, self.llm_primary_model, "failed") - raise ValueError("智能解读服务暂不可用，请稍后重试。") from last_error return generate() @@ -1867,53 +1840,35 @@ class DashboardService: if regime not in REGIMES: raise ValueError("市场阶段不支持。") notice = "" - compiled = None - primary_error = "" source = self.llm_source - started = datetime.now(timezone.utc) if source == "platform": - self.ensure_llm_access("screener") - if self.llm_configured: try: - compiled = compile_strategy_with_llm( - prompt, - regime, - self.llm_primary_api_key, - self.llm_primary_base_url, - self.llm_primary_model, + gateway_result = self.llm_gateway.call( + "screener", + "strategy-compiler-v1", + lambda profile: compile_strategy_with_llm( + prompt, + regime, + profile.api_key, + profile.base_url, + profile.model, + ), + (LLMCompilerError,), ) - except LLMCompilerError as exc: - primary_error = str(exc) - if compiled is None and self.llm_fallback_configured: - try: - compiled = compile_strategy_with_llm( - prompt, - regime, - self.llm_fallback_api_key, - self.llm_fallback_base_url, - self.llm_fallback_model, - ) - compiled["compiler"] = "llm_fallback" - notice = "智能策略生成服务已自动切换。" - except LLMCompilerError as exc: - fallback_error = str(exc) + compiled = gateway_result.value + if gateway_result.role == "fallback": + compiled["compiler"] = "llm_fallback" + notice = "智能策略生成服务已自动切换。" + except LLMGatewayError as exc: + if exc.code != "unavailable": + raise compiled = compile_local_strategy(prompt, regime) notice = "智能策略生成暂不可用，已使用本地模板。" - if compiled is None: + else: compiled = compile_local_strategy(prompt, regime) notice = "智能策略生成暂不可用，已使用本地模板。" compiled["formula"] = self.screener.validate_formula(compiled["formula"]) compiled["notice"] = notice - if source in {"personal", "platform"}: - elapsed = int((datetime.now(timezone.utc) - started).total_seconds() * 1000) - status = "success" if str(compiled.get("compiler") or "").startswith("llm") else "failed" - self.record_llm_usage( - "screener", - source, - str(compiled.get("model") or self.llm_primary_model), - status, - elapsed, - ) return compiled def save_screener_strategy(self, payload: dict[str, Any]) -> dict[str, Any]: @@ -2017,72 +1972,43 @@ class DashboardService: ) context = self._build_mentor_context(trade_date, question, skill) - source = self.ensure_llm_access("mentor") - profiles = [] - if self.llm_configured: - profiles.append( - ( - "primary", - self.llm_primary_api_key, - self.llm_primary_base_url, - self.llm_primary_model, - ) - ) - if self.llm_fallback_configured: - profiles.append( - ( - "fallback", - self.llm_fallback_api_key, - self.llm_fallback_base_url, - self.llm_fallback_model, - ) - ) - def generate(): - started = time.perf_counter() - last_error: Exception | None = None - for compiler, api_key, base_url, model in profiles: - try: - upstream = iter( - stream_with_mentor( - skill, context, question, history, api_key, base_url, model - ) - ) - first = next(upstream) - except (MentorAgentError, StopIteration) as exc: - last_error = exc - continue - answer_parts = [first] - yield {"type": "delta", "content": first} - try: - for chunk in upstream: - answer_parts.append(chunk) - yield {"type": "delta", "content": chunk} - except MentorAgentError as exc: - self.record_llm_usage("mentor", source, model, "failed") - raise ValueError("智能解读连接中断，请稍后重试。") from exc - answer = "".join(answer_parts).strip() - latency_ms = round((time.perf_counter() - started) * 1000) - self.database.save_mentor_exchange( - self.current_user_id, - mentor_id, - trade_date, + answer_parts: list[str] = [] + events = self.llm_gateway.stream( + "mentor", + f"mentor-skill-v1:{skill.skill_id}", + lambda profile: stream_with_mentor( + skill, + context, question, - answer, - context["data_trade_date"], - ) - self.record_llm_usage("mentor", source, model, "success", latency_ms) - yield { - "type": "meta", - "data_trade_date": context["data_trade_date"], - "notice": "智能解读已自动切换可用服务。" - if compiler == "fallback" - else "", - } - return - failed_model = profiles[-1][3] if profiles else self.llm_primary_model - self.record_llm_usage("mentor", source, failed_model, "failed") - raise ValueError("智能解读服务暂不可用，请稍后重试。") from last_error + history, + profile.api_key, + profile.base_url, + profile.model, + ), + (MentorAgentError,), + ) + for event in events: + if event.kind == "delta": + chunk = str(event.value or "") + answer_parts.append(chunk) + yield {"type": "delta", "content": chunk} + elif event.kind == "complete": + self.database.save_mentor_exchange( + self.current_user_id, + mentor_id, + trade_date, + question, + "".join(answer_parts).strip(), + context["data_trade_date"], + ) + yield { + "type": "meta", + "data_trade_date": context["data_trade_date"], + "notice": "智能解读已自动切换可用服务。" + if event.role == "fallback" + else "", + } return generate() @@ -3155,45 +3081,19 @@ class DashboardService: return bool(reading and str(reading.get("answer") or "").rstrip().endswith("……")) def _call_heaven_agent(self, mode: str, context: dict[str, Any]) -> tuple[dict[str, Any], str]: - source = self.ensure_llm_access(f"heaven_{mode}") - primary_error = "" - if self.llm_configured: - try: - result = interpret_heaven( - mode, - context, - self.llm_primary_api_key, - self.llm_primary_base_url, - self.llm_primary_model, - ) - self.record_llm_usage( - f"heaven_{mode}", source, str(result.get("model") or ""), - "success", int(result.get("latency_ms") or 0), - ) - return result, "primary" - except HeavenAgentError as exc: - primary_error = str(exc) - if self.llm_fallback_configured: - try: - result = interpret_heaven( - mode, - context, - self.llm_fallback_api_key, - self.llm_fallback_base_url, - self.llm_fallback_model, - ) - self.record_llm_usage( - f"heaven_{mode}", source, str(result.get("model") or ""), - "success", int(result.get("latency_ms") or 0), - ) - return result, "fallback" - except HeavenAgentError as exc: - self.record_llm_usage( - f"heaven_{mode}", source, self.llm_fallback_model, "failed" - ) - raise ValueError("智能解读服务暂不可用，请稍后重试。") from exc - self.record_llm_usage(f"heaven_{mode}", source, self.llm_primary_model, "failed") - raise ValueError("智能解读服务暂不可用，请稍后重试。") + result = self.llm_gateway.call( + f"heaven_{mode}", + f"heaven-{mode}-v1", + lambda profile: interpret_heaven( + mode, + context, + profile.api_key, + profile.base_url, + profile.model, + ), + (HeavenAgentError,), + ) + return result.value, result.role def _heaven_index_context( self, diff --git a/tests/test_database_migrations.py b/tests/test_database_migrations.py index 3b85ab8..efccab1 100644 --- a/tests/test_database_migrations.py +++ b/tests/test_database_migrations.py @@ -23,6 +23,7 @@ class DatabaseMigrationTests(unittest.TestCase): [ ("0001", "adopt_legacy_schema"), ("0002", "create_job_runs"), + ("0003", "extend_llm_audit"), ], ) ReviewDatabase(path) @@ -30,7 +31,7 @@ class DatabaseMigrationTests(unittest.TestCase): count = connection.execute( "SELECT COUNT(*) AS count FROM schema_migrations" ).fetchone()["count"] - self.assertEqual(count, 2) + self.assertEqual(count, 3) def test_connection_factory_enables_required_pragmas(self) -> None: with tempfile.TemporaryDirectory() as root: diff --git a/tests/test_llm_gateway.py b/tests/test_llm_gateway.py new file mode 100644 index 0000000..7e68769 --- /dev/null +++ b/tests/test_llm_gateway.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import unittest + +from backend.llm import LLMGateway, LLMGatewayError + + +class ProviderFailure(RuntimeError): + pass + + +class FakeDatabase: + def __init__(self, used: int = 0) -> None: + self.used = used + self.audit: list[dict[str, object]] = [] + + def count_llm_usage_since(self, user_id: int, source: str, since: str) -> int: + return self.used + + def record_llm_usage( + self, + user_id: int, + feature: str, + source: str, + model: str, + status: str, + latency_ms: int, + **metadata, + ) -> None: + self.audit.append( + { + "user_id": user_id, + "feature": feature, + "source": source, + "model": model, + "status": status, + **metadata, + } + ) + + +def profile() -> dict[str, object]: + return { + "source": "platform", + "primary": {"api_key": "p", "base_url": "https://p", "model": "primary"}, + "fallback": {"api_key": "f", "base_url": "https://f", "model": "fallback"}, + } + + +class LLMGatewayTests(unittest.TestCase): + def gateway(self, database: FakeDatabase | None = None) -> LLMGateway: + database = database or FakeDatabase() + return LLMGateway( + database=database, + user_id_supplier=lambda: 7, + membership_supplier=lambda: {"active": True}, + settings_supplier=lambda: {"member_daily_limit": 50}, + profile_supplier=profile, + ) + + def test_non_streaming_call_falls_back_and_audits_once(self) -> None: + database = FakeDatabase() + gateway = self.gateway(database) + + def invoke(model): + if model.role == "primary": + raise ProviderFailure("primary failed") + return {"answer": "ok"} + + result = gateway.call("heaven_trend", "heaven-trend-v1", invoke, (ProviderFailure,)) + + self.assertEqual(result.role, "fallback") + self.assertEqual(result.value, {"answer": "ok"}) + self.assertEqual(len(database.audit), 1) + self.assertEqual(database.audit[0]["model"], "fallback") + self.assertEqual(database.audit[0]["prompt_version"], "heaven-trend-v1") + + def test_stream_falls_back_before_first_delta(self) -> None: + database = FakeDatabase() + gateway = self.gateway(database) + + def invoke(model): + if model.role == "primary": + raise ProviderFailure("primary failed") + yield "a" + yield "b" + + events = list(gateway.stream("mentor", "mentor-v1", invoke, (ProviderFailure,))) + + self.assertEqual([event.value for event in events[:-1]], ["a", "b"]) + self.assertEqual(events[-1].kind, "complete") + self.assertEqual(events[-1].role, "fallback") + self.assertEqual(database.audit[0]["status"], "success") + + def test_stream_does_not_switch_model_after_output_started(self) -> None: + database = FakeDatabase() + gateway = self.gateway(database) + + def invoke(model): + yield "first" + raise ProviderFailure(f"{model.role} interrupted") + + iterator = gateway.stream("assistant", "assistant-v1", invoke, (ProviderFailure,)) + self.assertEqual(next(iterator).value, "first") + with self.assertRaisesRegex(LLMGatewayError, "连接中断"): + list(iterator) + self.assertEqual(len(database.audit), 1) + self.assertEqual(database.audit[0]["model"], "primary") + self.assertEqual(database.audit[0]["status"], "failed") + + def test_daily_quota_is_enforced_before_provider_call(self) -> None: + gateway = self.gateway(FakeDatabase(used=50)) + with self.assertRaisesRegex(LLMGatewayError, "额度已用完"): + gateway.call("mentor", "mentor-v1", lambda model: "unused", (ProviderFailure,)) + + +if __name__ == "__main__": + unittest.main()