From 4b451ef97ccaaddb1da4212853e56858d7434d7e Mon Sep 17 00:00:00 2001 From: wangwei Date: Thu, 2 Jul 2026 14:13:58 +0800 Subject: [PATCH] docs: add implementation plan for System Status AI model usage tracking Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../2026-07-02-status-llm-model-usage.md | 1755 +++++++++++++++++ 1 file changed, 1755 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-02-status-llm-model-usage.md diff --git a/docs/superpowers/plans/2026-07-02-status-llm-model-usage.md b/docs/superpowers/plans/2026-07-02-status-llm-model-usage.md new file mode 100644 index 0000000..26584ba --- /dev/null +++ b/docs/superpowers/plans/2026-07-02-status-llm-model-usage.md @@ -0,0 +1,1755 @@ +# System Status — AI Model Connections & Token Usage Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a new "AI Models" panel to the existing System Status page showing which LLM/Embedding/Reranker models are configured, their connection status, and cumulative token usage since process start. + +**Architecture:** Instrument the single shared LLM client factory (`get_llm_client()`) with a transparent `TrackedLLMClient` decorator so every call site (main answer generation, HyDE, Agentic RAG steps, compliance review, document summarization, regulation-perception analysis) is captured automatically. Instrument the embedding provider and reranker directly (each has only one implementation). All data is held in-memory by a new `ModelUsageTracker` singleton (`app/shared/model_usage_tracker.py`); two new endpoints on the existing `/status` router expose it; a new card on `StatusPage.tsx` displays it. + +**Tech Stack:** FastAPI, Python 3.10+ dataclasses, `asyncio.to_thread`/`asyncio.gather` (existing convention), React 19 + TypeScript, existing `fetchAPI` helper. + +**Design reference:** `docs/superpowers/specs/2026-07-02-status-llm-model-usage-design.md` — read this first for the "why" behind each decision below. + +## Global Constraints + +- Every new/modified Python file under `backend/` must have an English module docstring, docstrings on every class/function, and at least one meaningful `#` comment (project-wide rule; see `AGENTS.md`). +- No new business orchestration goes in `backend/app/services/*` or `backend/app/workflows/*` — new code goes in `backend/app/shared/` (cross-cutting support, same tier as `bootstrap.py`) or directly instruments the single existing infrastructure implementation. +- In-memory only — no new database table, no new dependency. Token counts reset on backend restart; this is an explicit, approved product decision, not a shortcut to fix later. +- Tracking code must never raise into a real request path — every public tracker method swallows its own exceptions and logs a warning. +- Frontend: desktop-first, no responsive/mobile work (`AGENTS.md`). No new frontend test framework — this repo has none (`frontend/package.json` has no test script); verify with `npm --prefix frontend run lint` and `npm --prefix frontend run build` only. +- Backend tests run via `pytest` (project already depends on `pytest>=7.0.0`, `pytest-asyncio>=0.21.0`). Root `pyproject.toml` sets `testpaths = ["tests"]`, so bare `pytest` only auto-discovers the root `tests/` directory — tests under `backend/tests/` must be invoked with an explicit path. Use `uv run pytest -v` (per `AGENTS.md`, this resolves within the project's `.venv`; if `uv`/`.venv` are not yet set up, run `./dev.sh setup` or `dev.bat setup` first, per `AGENTS.md`). +- `LLMFactory._global_instances` and `get_model_usage_tracker()` are both process-wide singletons/caches. Every test that touches them must clear their state in a fixture to avoid cross-test pollution — this is called out explicitly in the tasks below. + +--- + +### Task 1: `ModelUsageTracker` — in-memory usage/connection registry + +**Files:** +- Create: `backend/app/shared/model_usage_tracker.py` +- Create: `backend/tests/observability/__init__.py` +- Test: `backend/tests/observability/test_model_usage_tracker.py` + +**Interfaces:** +- Produces: `ModelUsageEntry` (dataclass: `provider: str`, `model: str`, `total_tokens: int`, `prompt_tokens: int`, `completion_tokens: int`, `call_count_ok: int`, `call_count_error: int`, `last_called_at: datetime | None`, `last_latency_ms: int | None`, `last_error: str | None`, property `status -> str` returning `"never_called" | "ok" | "error"`); `ModelUsageTracker` (methods `record(*, provider, model, success, usage=None, latency_ms=None, error=None) -> None`, `snapshot() -> dict[str, ModelUsageEntry]`, `get(provider, model) -> ModelUsageEntry | None`); `get_model_usage_tracker() -> ModelUsageTracker` (process-wide singleton via `@lru_cache`). + +- [ ] **Step 1: Write the failing test** + +Create `backend/tests/observability/__init__.py` (empty file, matches the `__init__.py`-per-test-package convention already used by `backend/tests/perception/` and `backend/tests/compliance/`). + +Create `backend/tests/observability/test_model_usage_tracker.py`: + +```python +"""Unit tests for ModelUsageTracker — no mocking needed, pure in-memory state.""" + +from __future__ import annotations + +from app.shared.model_usage_tracker import ModelUsageEntry, ModelUsageTracker, get_model_usage_tracker + + +def test_never_called_model_has_no_entry(): + """A tracker that has never recorded a call returns None from get().""" + tracker = ModelUsageTracker() + assert tracker.get("deepseek", "deepseek-v4-flash") is None + + +def test_record_success_accumulates_tokens_and_calls(): + """Two successful calls accumulate tokens and call_count_ok.""" + tracker = ModelUsageTracker() + tracker.record( + provider="deepseek", model="deepseek-v4-flash", success=True, + usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, latency_ms=100, + ) + tracker.record( + provider="deepseek", model="deepseek-v4-flash", success=True, + usage={"prompt_tokens": 20, "completion_tokens": 8, "total_tokens": 28}, latency_ms=200, + ) + entry = tracker.get("deepseek", "deepseek-v4-flash") + assert entry is not None + assert entry.total_tokens == 43 + assert entry.prompt_tokens == 30 + assert entry.completion_tokens == 13 + assert entry.call_count_ok == 2 + assert entry.call_count_error == 0 + assert entry.status == "ok" + assert entry.last_latency_ms == 200 + + +def test_record_error_sets_error_status_without_losing_prior_tokens(): + """A failed call after successful ones flips status to 'error' but keeps accumulated tokens.""" + tracker = ModelUsageTracker() + tracker.record(provider="qwen", model="qwen3.5-flash", success=True, usage={"total_tokens": 50}, latency_ms=50) + tracker.record(provider="qwen", model="qwen3.5-flash", success=False, error="HTTP 500", latency_ms=30) + entry = tracker.get("qwen", "qwen3.5-flash") + assert entry.total_tokens == 50 + assert entry.call_count_ok == 1 + assert entry.call_count_error == 1 + assert entry.status == "error" + assert entry.last_error == "HTTP 500" + + +def test_record_success_after_error_clears_last_error(): + """A later successful call clears last_error and status returns to 'ok'.""" + tracker = ModelUsageTracker() + tracker.record(provider="qwen", model="qwen3.5-flash", success=False, error="timeout", latency_ms=30) + tracker.record(provider="qwen", model="qwen3.5-flash", success=True, usage={"total_tokens": 5}, latency_ms=40) + entry = tracker.get("qwen", "qwen3.5-flash") + assert entry.status == "ok" + assert entry.last_error is None + + +def test_record_never_raises_on_bad_usage_dict(): + """A malformed usage value (wrong type) is swallowed, not raised, and does not corrupt other entries.""" + tracker = ModelUsageTracker() + tracker.record(provider="embedding", model="text-embedding-v3", success=True, usage="not-a-dict", latency_ms=10) # type: ignore[arg-type] + # Must not raise, and must not have created a corrupted entry that breaks snapshot(). + snapshot = tracker.snapshot() + assert isinstance(snapshot, dict) + + +def test_snapshot_returns_independent_copy(): + """snapshot() returns a dict that can be safely mutated without affecting the tracker.""" + tracker = ModelUsageTracker() + tracker.record(provider="deepseek", model="deepseek-v4-flash", success=True, usage={"total_tokens": 1}, latency_ms=1) + snap = tracker.snapshot() + snap.clear() + assert tracker.get("deepseek", "deepseek-v4-flash") is not None + + +def test_get_model_usage_tracker_returns_singleton(): + """get_model_usage_tracker() always returns the same process-wide instance.""" + assert get_model_usage_tracker() is get_model_usage_tracker() +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest backend/tests/observability/test_model_usage_tracker.py -v` +Expected: FAIL — `ModuleNotFoundError: No module named 'app.shared.model_usage_tracker'` + +- [ ] **Step 3: Write minimal implementation** + +Create `backend/app/shared/model_usage_tracker.py`: + +```python +"""In-memory registry that tracks per-model call outcomes and token usage. + +This module lives in `app/shared` — the same cross-cutting-support tier as +`bootstrap.py` — because it is not business logic: it exists purely so the +System Status page can show which AI models (main LLM, HyDE LLM, embedding, +reranker) are configured, whether their most recent call succeeded, and how +many tokens they have consumed since this process started. Tracking here +must never disrupt a real user-facing call: every public method swallows its +own exceptions and logs a warning instead of raising. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from datetime import datetime, timezone +from functools import lru_cache + +from loguru import logger + + +@dataclass +class ModelUsageEntry: + """Represent accumulated usage/connection state for one provider+model pair.""" + + provider: str + model: str + total_tokens: int = 0 + prompt_tokens: int = 0 + completion_tokens: int = 0 + call_count_ok: int = 0 + call_count_error: int = 0 + last_called_at: datetime | None = None + last_latency_ms: int | None = None + last_error: str | None = None + + @property + def status(self) -> str: + """Derive never_called/ok/error from call history. + + The "disabled" status (reranker only, when turned off in settings) is + NOT decided here: this dataclass has no access to live settings. The + API route layer (Task 6) applies that override on top of this value, + so config always wins over stale historical data. + """ + if self.last_called_at is None: + return "never_called" + return "error" if self.last_error else "ok" + + +class ModelUsageTracker: + """Thread-safe in-memory registry of per-model call/usage stats. + + Keyed by "{provider}:{model}" rather than by business role (main LLM / + HyDE / embedding / reranker) so that any future call site is captured + automatically, even before anyone teaches this class about its role. + """ + + def __init__(self) -> None: + """Initialize an empty registry guarded by a single lock.""" + self._entries: dict[str, ModelUsageEntry] = {} + # One coarse lock is enough: record() runs at most a few times per + # request, and snapshot() is only read by the low-traffic status page. + self._lock = threading.Lock() + + def record( + self, + *, + provider: str, + model: str, + success: bool, + usage: dict | None = None, + latency_ms: int | None = None, + error: str | None = None, + ) -> None: + """Record the outcome of one call to provider/model. + + Never raises: any internal failure is logged and swallowed so a bug + in observability code cannot break a real LLM/embedding/reranker call. + """ + try: + key = f"{provider}:{model}" + usage = usage if isinstance(usage, dict) else {} + with self._lock: + entry = self._entries.setdefault(key, ModelUsageEntry(provider=provider, model=model)) + entry.total_tokens += int(usage.get("total_tokens", 0) or 0) + entry.prompt_tokens += int(usage.get("prompt_tokens", 0) or 0) + entry.completion_tokens += int(usage.get("completion_tokens", 0) or 0) + if success: + entry.call_count_ok += 1 + entry.last_error = None + else: + entry.call_count_error += 1 + entry.last_error = error or "unknown error" + entry.last_called_at = datetime.now(timezone.utc) + entry.last_latency_ms = latency_ms + except Exception as exc: # noqa: BLE001 - tracking must never break a real call + logger.warning("ModelUsageTracker.record failed for {}:{} - {}", provider, model, exc) + + def snapshot(self) -> dict[str, ModelUsageEntry]: + """Return a shallow copy of all tracked entries, safe to mutate by the caller.""" + with self._lock: + return dict(self._entries) + + def get(self, provider: str, model: str) -> ModelUsageEntry | None: + """Return the entry for one provider/model pair, or None if never recorded.""" + return self.snapshot().get(f"{provider}:{model}") + + +@lru_cache +def get_model_usage_tracker() -> ModelUsageTracker: + """Return the process-wide singleton tracker (mirrors get_settings()/get_llm_factory()).""" + return ModelUsageTracker() +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest backend/tests/observability/test_model_usage_tracker.py -v` +Expected: PASS (7 tests) + +- [ ] **Step 5: Commit** + +```bash +git add backend/app/shared/model_usage_tracker.py backend/tests/observability/ +git commit -m "feat: add ModelUsageTracker for per-model token/connection tracking" +``` + +--- + +### Task 2: `TrackedLLMClient` — transparent usage-recording decorator + +**Files:** +- Create: `backend/app/services/llm/tracked_client.py` +- Create: `backend/tests/observability/test_tracked_client.py` + +**Interfaces:** +- Consumes: `ModelUsageTracker` (Task 1, `backend/app/shared/model_usage_tracker.py`), `BaseLLMClient`/`LLMResponse` (`backend/app/services/llm/base_client.py`, existing). +- Produces: `TrackedLLMClient(inner: BaseLLMClient, tracker: ModelUsageTracker)` with `.chat(...)`, `.stream_chat(...)`, and `__getattr__` passthrough for everything else (`get_available_models`, `close`, `.config`). + +- [ ] **Step 1: Write the failing test** + +Create `backend/tests/observability/test_tracked_client.py`: + +```python +"""Unit tests for TrackedLLMClient — verifies transparent delegation + recording.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from app.services.llm.base_client import LLMConfig, LLMProvider, LLMResponse +from app.services.llm.tracked_client import TrackedLLMClient +from app.shared.model_usage_tracker import ModelUsageTracker + + +def _make_inner(model: str = "deepseek-v4-flash") -> MagicMock: + """Build a MagicMock standing in for a concrete BaseLLMClient subclass.""" + inner = MagicMock() + inner.config = LLMConfig( + provider=LLMProvider.DEEPSEEK, model=model, api_key="test-key", base_url="http://example.test/v1", + ) + return inner + + +def test_chat_delegates_and_returns_unchanged_response(): + """chat() must return exactly what the wrapped client returned.""" + inner = _make_inner() + expected = LLMResponse(content="hello", model="deepseek-v4-flash", usage={"total_tokens": 12}) + inner.chat.return_value = expected + tracker = ModelUsageTracker() + + tracked = TrackedLLMClient(inner, tracker) + result = tracked.chat([{"role": "user", "content": "hi"}]) + + assert result is expected + inner.chat.assert_called_once_with([{"role": "user", "content": "hi"}], None, None, None) + + +def test_chat_records_success_and_tokens(): + """A successful chat() call must be recorded under 'deepseek:deepseek-v4-flash'.""" + inner = _make_inner() + inner.chat.return_value = LLMResponse(content="hi", model="deepseek-v4-flash", usage={"total_tokens": 42}) + tracker = ModelUsageTracker() + + TrackedLLMClient(inner, tracker).chat([{"role": "user", "content": "hi"}]) + + entry = tracker.get("deepseek", "deepseek-v4-flash") + assert entry is not None + assert entry.total_tokens == 42 + assert entry.status == "ok" + + +def test_chat_records_error_from_response(): + """A chat() call that returns an error-carrying LLMResponse is recorded as a failure.""" + inner = _make_inner() + inner.chat.return_value = LLMResponse(content="", model="deepseek-v4-flash", error="API error: 500") + tracker = ModelUsageTracker() + + TrackedLLMClient(inner, tracker).chat([{"role": "user", "content": "hi"}]) + + entry = tracker.get("deepseek", "deepseek-v4-flash") + assert entry.status == "error" + assert entry.last_error == "API error: 500" + + +def test_getattr_forwards_to_inner_client(): + """Attributes not defined on TrackedLLMClient must forward to the wrapped client.""" + inner = _make_inner() + inner.get_available_models.return_value = ["deepseek-v4-flash"] + tracked = TrackedLLMClient(inner, ModelUsageTracker()) + + assert tracked.get_available_models() == ["deepseek-v4-flash"] + assert tracked.config is inner.config + + +def test_stream_chat_records_call_without_token_usage(): + """stream_chat() must record a call (latency/success) but not fabricate token counts.""" + inner = _make_inner() + inner.stream_chat.return_value = iter(["chunk-1", "chunk-2"]) + tracker = ModelUsageTracker() + + chunks = list(TrackedLLMClient(inner, tracker).stream_chat([{"role": "user", "content": "hi"}])) + + assert chunks == ["chunk-1", "chunk-2"] + entry = tracker.get("deepseek", "deepseek-v4-flash") + assert entry.call_count_ok == 1 + assert entry.total_tokens == 0 +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest backend/tests/observability/test_tracked_client.py -v` +Expected: FAIL — `ModuleNotFoundError: No module named 'app.services.llm.tracked_client'` + +- [ ] **Step 3: Write minimal implementation** + +Create `backend/app/services/llm/tracked_client.py`: + +```python +"""Transparent decorator around BaseLLMClient implementations. + +Records per-call token usage, latency, and success/failure into a +ModelUsageTracker without changing any caller-visible behavior. +""" + +from __future__ import annotations + +import time +from typing import Any, Dict, List, Optional + +from app.shared.model_usage_tracker import ModelUsageTracker + +from .base_client import BaseLLMClient, LLMResponse +from .tool_types import Tool + + +class TrackedLLMClient: + """Wrap any BaseLLMClient and record its usage into a ModelUsageTracker. + + Deliberately does NOT subclass BaseLLMClient: that ABC declares abstract + methods (_init_client, get_available_models) with no meaningful override + here, and subclassing would make Python refuse to instantiate this class + ("Can't instantiate abstract class") before __getattr__ ever got a chance + to forward the call. Plain composition + __getattr__ delegation works + because every caller in this codebase only ever uses duck-typed access: + .chat(), .stream_chat(), .get_available_models(), .close(), .config. + """ + + def __init__(self, inner: BaseLLMClient, tracker: ModelUsageTracker) -> None: + """Store the wrapped client and the tracker to report into.""" + self._inner = inner + self._tracker = tracker + + def chat( + self, + messages: List[Dict[str, str]], + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + tools: Optional[List[Tool]] = None, + **kwargs: Any, + ) -> LLMResponse: + """Delegate to the wrapped client's chat(), then record the outcome.""" + start = time.time() + response = self._inner.chat(messages, max_tokens, temperature, tools, **kwargs) + # Key by the *configured* model, not response.model, so lookups driven + # by settings (llm_model / hyde_llm_model) always match what we recorded. + self._tracker.record( + provider=self._inner.config.provider.value, + model=self._inner.config.model, + success=response.is_success, + usage=response.usage, + latency_ms=int((time.time() - start) * 1000), + error=response.error, + ) + return response + + def stream_chat(self, messages: List[Dict[str, str]], *args: Any, **kwargs: Any): + """Delegate to the wrapped client's stream_chat(), recording call outcome only. + + Token usage is NOT recorded here: none of the current provider + stream_chat() implementations parse a trailing usage chunk from the + gateway (see the design doc's Known Limitations), so accumulating a + token count here would silently be wrong. Only call success/failure + and latency are tracked for streaming calls. + """ + start = time.time() + error: Optional[str] = None + try: + for chunk in self._inner.stream_chat(messages, *args, **kwargs): + yield chunk + except Exception as exc: # noqa: BLE001 - report, then re-raise unchanged + error = str(exc) + raise + finally: + self._tracker.record( + provider=self._inner.config.provider.value, + model=self._inner.config.model, + success=error is None, + latency_ms=int((time.time() - start) * 1000), + error=error, + ) + + def __getattr__(self, name: str) -> Any: + """Forward any other attribute/method access to the wrapped client.""" + return getattr(self._inner, name) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest backend/tests/observability/test_tracked_client.py -v` +Expected: PASS (5 tests) + +- [ ] **Step 5: Commit** + +```bash +git add backend/app/services/llm/tracked_client.py backend/tests/observability/test_tracked_client.py +git commit -m "feat: add TrackedLLMClient decorator for transparent usage recording" +``` + +--- + +### Task 3: Wire `TrackedLLMClient` into `LLMFactory` + +**Files:** +- Modify: `backend/app/services/llm/llm_factory.py:1-30` (imports), `:39-83` (`create()`), `:140-145` (`get_cached()`), `:199-212` (`get_llm_client()`) +- Test: `backend/tests/observability/test_llm_factory_tracking.py` + +**Interfaces:** +- Consumes: `TrackedLLMClient` (Task 2), `get_model_usage_tracker()` (Task 1). +- Produces: `get_llm_client(...)` now returns a `TrackedLLMClient`-wrapped instance (return type becomes `BaseLLMClient | TrackedLLMClient`) — no signature change for any caller. + +- [ ] **Step 1: Write the failing test** + +Create `backend/tests/observability/test_llm_factory_tracking.py`: + +```python +"""Verifies get_llm_client() returns a usage-tracked client end to end.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from app.services.llm.llm_factory import LLMFactory, get_llm_client +from app.services.llm.tracked_client import TrackedLLMClient +from app.shared.model_usage_tracker import get_model_usage_tracker + + +@pytest.fixture(autouse=True) +def _reset_singletons(): + """Clear the two process-wide singletons this test touches, before and after. + + LLMFactory._global_instances and get_model_usage_tracker() both persist + for the life of the process; without this fixture, tests would leak + cached clients/usage data into each other and become order-dependent. + """ + LLMFactory._global_instances.clear() + get_model_usage_tracker().snapshot() # no-op read, just documents intent + get_model_usage_tracker()._entries.clear() + yield + LLMFactory._global_instances.clear() + get_model_usage_tracker()._entries.clear() + + +def test_get_llm_client_returns_tracked_client(): + """get_llm_client() must return a TrackedLLMClient, not the raw provider client.""" + with patch("app.services.llm.llm_factory.DeepSeekClient") as mock_cls: + mock_cls.return_value = MagicMock() + client = get_llm_client(provider="deepseek", model="deepseek-v4-flash", api_key="test-key") + assert isinstance(client, TrackedLLMClient) + + +def test_get_llm_client_caches_the_tracked_instance(): + """A second call with the same provider/model must return the same TrackedLLMClient.""" + with patch("app.services.llm.llm_factory.DeepSeekClient") as mock_cls: + mock_cls.return_value = MagicMock() + first = get_llm_client(provider="deepseek", model="deepseek-v4-flash", api_key="test-key") + second = get_llm_client(provider="deepseek", model="deepseek-v4-flash", api_key="test-key") + assert first is second +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest backend/tests/observability/test_llm_factory_tracking.py -v` +Expected: FAIL — `assert isinstance(client, TrackedLLMClient)` fails because `create()` does not wrap yet. + +- [ ] **Step 3: Write minimal implementation** + +In `backend/app/services/llm/llm_factory.py`, add two imports near the top (after the existing `from .qwen_client import QwenClient, QwenVLClient` line): + +```python +from .tracked_client import TrackedLLMClient +from app.shared.model_usage_tracker import get_model_usage_tracker +``` + +Modify `create()` — insert the wrapping step between building the client and caching it: + +```python + # Keep provider-specific behavior explicit so debugging stays straightforward. + client = self._create_client(config) + + # Wrap in TrackedLLMClient so every call site (agentic, HyDE, perception, + # compliance, document summarization, main answer generation) is recorded + # without each of them needing to know about usage tracking. + tracked_client = TrackedLLMClient(client, get_model_usage_tracker()) + + # Keep provider-specific behavior explicit so debugging stays straightforward. + LLMFactory._global_instances[cache_key] = tracked_client + + logger.info(f"LLM客户端创建成功并缓存: {provider} - {model}") + return tracked_client +``` + +Update the `create()` method's return type annotation from `-> BaseLLMClient` to `-> "BaseLLMClient | TrackedLLMClient"`. + +Update `get_cached()`'s return type annotation from `Optional[BaseLLMClient]` to `"BaseLLMClient | TrackedLLMClient | None"`. + +Update the module-level `get_llm_client()` function's return type annotation from `-> BaseLLMClient` to `-> "BaseLLMClient | TrackedLLMClient"`. + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest backend/tests/observability/test_llm_factory_tracking.py -v` +Expected: PASS (2 tests) + +Then re-run Tasks 1-2's tests together to confirm no regressions: + +Run: `uv run pytest backend/tests/observability/ -v` +Expected: PASS (all tests from Tasks 1-3) + +- [ ] **Step 5: Commit** + +```bash +git add backend/app/services/llm/llm_factory.py backend/tests/observability/test_llm_factory_tracking.py +git commit -m "feat: wrap LLM clients with TrackedLLMClient in LLMFactory" +``` + +--- + +### Task 4: Instrument the embedding provider + +**Files:** +- Modify: `backend/app/infrastructure/embedding/openai_compatible_embedding_provider.py` (full file, 79 lines) +- Test: `backend/tests/observability/test_embedding_usage_tracking.py` + +**Interfaces:** +- Consumes: `get_model_usage_tracker()` (Task 1). +- Produces: no interface change — `OpenAICompatibleEmbeddingProvider.embed_texts`/`embed_query` behave identically; only a side effect (tracking) is added. Tracker key is `"embedding:{settings.embedding_model}"`. + +- [ ] **Step 1: Write the failing test** + +Create `backend/tests/observability/test_embedding_usage_tracking.py`: + +```python +"""Verifies OpenAICompatibleEmbeddingProvider records usage into ModelUsageTracker.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from app.infrastructure.embedding.openai_compatible_embedding_provider import OpenAICompatibleEmbeddingProvider +from app.shared.model_usage_tracker import get_model_usage_tracker + + +@pytest.fixture(autouse=True) +def _reset_tracker(): + """Clear the process-wide tracker before and after each test in this file.""" + get_model_usage_tracker()._entries.clear() + yield + get_model_usage_tracker()._entries.clear() + + +def _fake_response(usage: dict) -> MagicMock: + """Build a fake httpx.Response-like object for a successful embeddings call.""" + resp = MagicMock(spec=httpx.Response) + resp.raise_for_status.return_value = None + resp.json.return_value = { + "data": [{"index": 0, "embedding": [0.1] * 1024}], + "usage": usage, + } + return resp + + +def test_successful_embed_records_usage(): + """A successful embeddings call must record token usage under 'embedding:'.""" + provider = OpenAICompatibleEmbeddingProvider() + provider.api_key = "test-key" + with patch("httpx.post", return_value=_fake_response({"prompt_tokens": 3, "total_tokens": 3})): + provider.embed_query("hello") + + entry = get_model_usage_tracker().get("embedding", provider.model) + assert entry is not None + assert entry.total_tokens == 3 + assert entry.status == "ok" + + +def test_failed_embed_records_error(): + """An HTTP error from the embeddings endpoint must be recorded as a failure, then re-raised.""" + provider = OpenAICompatibleEmbeddingProvider() + provider.api_key = "test-key" + failing_response = MagicMock(spec=httpx.Response) + failing_response.raise_for_status.side_effect = httpx.HTTPStatusError( + "boom", request=MagicMock(), response=MagicMock(status_code=500, text="boom") + ) + failing_response.text = "boom" + with patch("httpx.post", return_value=failing_response): + with pytest.raises(httpx.HTTPStatusError): + provider.embed_query("hello") + + entry = get_model_usage_tracker().get("embedding", provider.model) + assert entry is not None + assert entry.status == "error" + assert entry.call_count_error == 1 +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest backend/tests/observability/test_embedding_usage_tracking.py -v` +Expected: FAIL — `entry` is `None` (no tracking wired yet). + +- [ ] **Step 3: Write minimal implementation** + +Replace the full contents of `backend/app/infrastructure/embedding/openai_compatible_embedding_provider.py`: + +```python +"""Implement infrastructure support for openai compatible embedding provider.""" + +from __future__ import annotations + +import os +import time + +import httpx + +from app.config.settings import settings +from app.domain.retrieval import EmbeddingProvider +from app.shared.model_usage_tracker import get_model_usage_tracker +# Keep adapter behavior explicit so integration details remain easy to audit. + +EMBEDDING_BATCH_SIZE = 8 + + + +class OpenAICompatibleEmbeddingProvider(EmbeddingProvider): + """Provide the Open A I Compatible Embedding Provider provider.""" + def __init__(self) -> None: + """Initialize the Open A I Compatible Embedding Provider instance.""" + self.base_url = settings.embedding_base_url.rstrip("/") + self.api_key = ( + settings.embedding_api_key + or os.getenv("OPENAI_API_KEY", "") + or os.getenv("QWEN_API_KEY", "") + or os.getenv("DEEPSEEK_API_KEY", "") + ) + self.model = settings.embedding_model + self.timeout = settings.embedding_timeout_seconds + self.dimension = settings.embedding_dim + + def _raise_for_status(self, response: httpx.Response, *, batch_size: int) -> None: + """Raise a detailed error so upstream gateway failures are easier to diagnose.""" + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + response_preview = response.text[:500].strip() + detail = ( + f"Embedding request failed for model={self.model}, batch_size={batch_size}, " + f"status={response.status_code}, url={response.request.url}, response={response_preview}" + ) + raise httpx.HTTPStatusError(detail, request=exc.request, response=exc.response) from exc + + def _request(self, texts: list[str]) -> list[list[float]]: + """Handle request for this module for the Open A I Compatible Embedding Provider instance.""" + if not self.api_key: + raise ValueError("缺少 EMBEDDING_API_KEY / OPENAI_API_KEY") + start = time.time() + try: + response = httpx.post( + f"{self.base_url}/embeddings", + headers={ + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + json={"model": self.model, "input": texts}, + timeout=self.timeout, + ) + self._raise_for_status(response, batch_size=len(texts)) + data = response.json() + except Exception as exc: + # Record the failed call so the Status page can show it as an error, + # then re-raise unchanged so existing callers keep their current behavior. + get_model_usage_tracker().record( + provider="embedding", + model=self.model, + success=False, + latency_ms=int((time.time() - start) * 1000), + error=str(exc), + ) + raise + vectors = [item["embedding"] for item in sorted(data.get("data", []), key=lambda item: item["index"])] + if any(len(vector) != self.dimension for vector in vectors): + raise ValueError(f"embedding 维度不匹配，期望 {self.dimension}") + # Record token usage from the OpenAI-compatible response, e.g. {"total_tokens": N}. + get_model_usage_tracker().record( + provider="embedding", + model=self.model, + success=True, + usage=data.get("usage", {}), + latency_ms=int((time.time() - start) * 1000), + ) + return vectors + + def embed_texts(self, texts: list[str]) -> list[list[float]]: + """Embed texts for the Open A I Compatible Embedding Provider instance.""" + if not texts: + return [] + vectors: list[list[float]] = [] + # Batch requests conservatively because some gateways reject larger embedding payloads. + for start in range(0, len(texts), EMBEDDING_BATCH_SIZE): + batch = texts[start:start + EMBEDDING_BATCH_SIZE] + vectors.extend(self._request(batch)) + return vectors + + def embed_query(self, text: str) -> list[float]: + """Embed query for the Open A I Compatible Embedding Provider instance.""" + vectors = self._request([text]) + return vectors[0] +``` + +(Only the `_request` method changed: added `import time`, `from app.shared.model_usage_tracker import get_model_usage_tracker`, wrapped the request in try/except to record failures, and record success with `usage` after the dimension check.) + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest backend/tests/observability/test_embedding_usage_tracking.py -v` +Expected: PASS (2 tests) + +- [ ] **Step 5: Commit** + +```bash +git add backend/app/infrastructure/embedding/openai_compatible_embedding_provider.py backend/tests/observability/test_embedding_usage_tracking.py +git commit -m "feat: record embedding call usage into ModelUsageTracker" +``` + +--- + +### Task 5: Instrument the reranker + +**Files:** +- Modify: `backend/app/infrastructure/vectorstore/cross_encoder_reranker.py:1-54` (imports + `rerank()`) +- Test: `backend/tests/observability/test_reranker_usage_tracking.py` + +**Interfaces:** +- Consumes: `get_model_usage_tracker()` (Task 1). +- Produces: no interface change — `OpenAICompatibleReranker.rerank()` behaves identically; tracker key is `"reranker:{settings.reranker_model}"`. No token usage is ever recorded for this role (TEI/Cohere-style rerank APIs do not return one) — only call success/failure and latency. + +- [ ] **Step 1: Write the failing test** + +Create `backend/tests/observability/test_reranker_usage_tracking.py`: + +```python +"""Verifies OpenAICompatibleReranker records call outcome (no tokens) into ModelUsageTracker.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from app.domain.retrieval import RetrievedChunk +from app.infrastructure.vectorstore.cross_encoder_reranker import OpenAICompatibleReranker +from app.shared.model_usage_tracker import get_model_usage_tracker + + +@pytest.fixture(autouse=True) +def _reset_tracker(): + """Clear the process-wide tracker before and after each test in this file.""" + get_model_usage_tracker()._entries.clear() + yield + get_model_usage_tracker()._entries.clear() + + +def _chunk(chunk_id: str, text: str) -> RetrievedChunk: + """Build a minimal RetrievedChunk for reranker tests.""" + return RetrievedChunk(chunk_id=chunk_id, doc_id="doc-1", doc_title="Doc", text=text, score=0.0) + + +def test_successful_rerank_records_call_without_tokens(): + """A successful rerank() call is recorded with call_count_ok but zero tokens.""" + reranker = OpenAICompatibleReranker(base_url="http://example.test", model="bge-reranker-v2-m3") + with patch.object(reranker, "_call_reranker", return_value=[0.9, 0.1]): + result = reranker.rerank("query", [_chunk("c1", "a"), _chunk("c2", "b")], top_k=2) + + assert len(result) == 2 + entry = get_model_usage_tracker().get("reranker", "bge-reranker-v2-m3") + assert entry is not None + assert entry.call_count_ok == 1 + assert entry.total_tokens == 0 + + +def test_failed_rerank_records_error_and_falls_back(): + """A rerank() call that raises internally is recorded as an error but still returns a fallback list.""" + reranker = OpenAICompatibleReranker(base_url="http://example.test", model="bge-reranker-v2-m3") + with patch.object(reranker, "_call_reranker", side_effect=RuntimeError("gateway down")): + result = reranker.rerank("query", [_chunk("c1", "a")], top_k=1) + + assert len(result) == 1 # existing fallback behavior: original order, unscored + entry = get_model_usage_tracker().get("reranker", "bge-reranker-v2-m3") + assert entry is not None + assert entry.call_count_error == 1 + assert entry.status == "error" +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest backend/tests/observability/test_reranker_usage_tracking.py -v` +Expected: FAIL — `entry` is `None` (no tracking wired yet). + +- [ ] **Step 3: Write minimal implementation** + +In `backend/app/infrastructure/vectorstore/cross_encoder_reranker.py`, add an import after the existing `from app.domain.retrieval import Reranker, RetrievedChunk` line: + +```python +from app.shared.model_usage_tracker import get_model_usage_tracker +``` + +Replace the `rerank()` method body with: + +```python + def rerank(self, query: str, chunks: list[RetrievedChunk], top_k: int) -> list[RetrievedChunk]: + """Return up to top_k chunks re-sorted by cross-encoder score.""" + if not chunks: + return [] + + texts = [chunk.text for chunk in chunks] + start = time.time() + try: + scores = self._call_reranker(query, texts) + except Exception as exc: + logger.warning("Reranker call failed ({}), falling back to original order: {}", type(exc).__name__, exc) + # Record the failure so the Status page reflects real reranker health. + get_model_usage_tracker().record( + provider="reranker", + model=self._model, + success=False, + latency_ms=int((time.time() - start) * 1000), + error=str(exc), + ) + return chunks[:top_k] + + elapsed_ms = int((time.time() - start) * 1000) + logger.debug("Reranker scored {} chunks in {}ms", len(chunks), elapsed_ms) + # TEI/Cohere-style rerank responses carry no token usage field — + # only call success/latency is meaningful for this role. + get_model_usage_tracker().record( + provider="reranker", + model=self._model, + success=True, + latency_ms=elapsed_ms, + ) + + ranked = sorted( + [(score, chunk) for score, chunk in zip(scores, chunks)], + key=lambda x: x[0], + reverse=True, + ) + result = [] + for score, chunk in ranked[:top_k]: + chunk.score = float(score) + result.append(chunk) + return result +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest backend/tests/observability/test_reranker_usage_tracking.py -v` +Expected: PASS (2 tests) + +Then run the full new test package together: + +Run: `uv run pytest backend/tests/observability/ -v` +Expected: PASS (all tests from Tasks 1-5) + +- [ ] **Step 5: Commit** + +```bash +git add backend/app/infrastructure/vectorstore/cross_encoder_reranker.py backend/tests/observability/test_reranker_usage_tracking.py +git commit -m "feat: record reranker call outcome into ModelUsageTracker" +``` + +--- + +### Task 6: `GET /status/models` and `POST /status/models/ping` routes + +**Files:** +- Modify: `backend/app/api/routes/status.py` (full file, 114 lines) +- Test: `tests/test_status_models_routes.py` (root `tests/`, matching `tests/test_auth_routes.py` convention for route-level tests) + +**Interfaces:** +- Consumes: `get_model_usage_tracker()` (Task 1), `get_llm_client()` (Task 3, now tracked), `get_embedding_provider()`/`get_reranker()` (existing bootstrap accessors, Task 4/5 now instrument what they return), `RetrievedChunk` (`app.domain.retrieval`, existing). +- Produces: `GET /api/v1/status/models` → `{"models": [<4 entries>]}`; `POST /api/v1/status/models/ping` → same shape, refreshed. Each entry: `{role, role_label, provider, model, enabled, status, total_tokens, call_count_ok, call_count_error, last_called_at, last_latency_ms, last_error, shares_usage_with}`. + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_status_models_routes.py`: + +```python +"""Integration tests for the /status/models routes. + +Uses FastAPI TestClient with mocked LLM/embedding/reranker clients so no +external gateway or database is required. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +from app.services.llm.base_client import LLMResponse +from app.shared.model_usage_tracker import get_model_usage_tracker + + +@pytest.fixture(autouse=True) +def _reset_tracker(): + """Clear the process-wide tracker before and after each test in this file.""" + get_model_usage_tracker()._entries.clear() + yield + get_model_usage_tracker()._entries.clear() + + +@pytest.fixture +def client(): + """Return a TestClient for the real app (status routes require no auth).""" + from app.api.main import app + with TestClient(app, raise_server_exceptions=False) as c: + yield c + + +def test_get_models_returns_four_roles_never_called_by_default(client): + """With no calls made yet, all 4 roles are returned with status 'never_called' or 'disabled'.""" + resp = client.get("/api/v1/status/models") + assert resp.status_code == 200 + body = resp.json() + roles = {m["role"] for m in body["models"]} + assert roles == {"main_llm", "hyde_llm", "embedding", "reranker"} + reranker_row = next(m for m in body["models"] if m["role"] == "reranker") + # Default .env.example ships RERANKER_ENABLED=false. + from app.config.settings import settings + assert reranker_row["enabled"] == settings.reranker_enabled + if not settings.reranker_enabled: + assert reranker_row["status"] == "disabled" + + +def test_get_models_reflects_recorded_usage(client): + """A previously recorded call must show up in total_tokens/status.""" + from app.config.settings import settings + get_model_usage_tracker().record( + provider=settings.llm_provider, model=settings.llm_model, success=True, usage={"total_tokens": 99}, + ) + resp = client.get("/api/v1/status/models") + main_row = next(m for m in resp.json()["models"] if m["role"] == "main_llm") + assert main_row["total_tokens"] == 99 + assert main_row["status"] == "ok" + + +def test_ping_models_calls_each_enabled_model_once(client): + """POST /status/models/ping must invoke chat()/embed_query() and return fresh statuses.""" + mock_llm_response = LLMResponse(content="pong", model="test-model", usage={"total_tokens": 1}) + mock_llm_client = MagicMock() + mock_llm_client.chat.return_value = mock_llm_response + mock_embedding = MagicMock() + mock_embedding.embed_query.return_value = [0.1] + + with patch("app.api.routes.status.get_llm_client", return_value=mock_llm_client), \ + patch("app.api.routes.status.get_embedding_provider", return_value=mock_embedding), \ + patch("app.api.routes.status.get_reranker", return_value=None): + resp = client.post("/api/v1/status/models/ping") + + assert resp.status_code == 200 + body = resp.json() + assert len(body["models"]) == 4 + assert mock_llm_client.chat.call_count >= 1 + mock_embedding.embed_query.assert_called_once() + + +def test_ping_models_survives_one_model_failing(client): + """If the LLM ping raises, embedding/reranker pings must still be attempted and a 200 returned.""" + mock_embedding = MagicMock() + mock_embedding.embed_query.return_value = [0.1] + + with patch("app.api.routes.status.get_llm_client", side_effect=RuntimeError("gateway down")), \ + patch("app.api.routes.status.get_embedding_provider", return_value=mock_embedding), \ + patch("app.api.routes.status.get_reranker", return_value=None): + resp = client.post("/api/v1/status/models/ping") + + assert resp.status_code == 200 + mock_embedding.embed_query.assert_called_once() +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/test_status_models_routes.py -v` +Expected: FAIL — `404 Not Found` for `/api/v1/status/models` (route does not exist yet). + +- [ ] **Step 3: Write minimal implementation** + +Replace the full contents of `backend/app/api/routes/status.py`: + +```python +"""Define API routes for status.""" + +import asyncio +import time +from typing import Any + +from fastapi import APIRouter + +from app.config.settings import settings +from app.domain.retrieval import RetrievedChunk +from app.services.llm.llm_factory import get_llm_client +from app.shared.bootstrap import ( + get_bm25_retriever, + get_binary_store, + get_conversation_store, + get_document_query_service, + get_embedding_provider, + get_reranker, + get_vector_index, +) +from app.shared.model_usage_tracker import get_model_usage_tracker + +router = APIRouter(prefix="/status", tags=["系统状态"]) + +# --------------------------------------------------------------------------- +# Simple TTL cache for /stats (avoids O(N) doc scan on every request) +# --------------------------------------------------------------------------- +_stats_cache: dict[str, Any] = {} +_stats_cache_time: float = 0.0 +_STATS_TTL_SECONDS: float = 10.0 + +# --------------------------------------------------------------------------- +# AI model roles surfaced on the Status page (Task: System Status AI models) +# --------------------------------------------------------------------------- +_MODEL_ROLES: dict[str, str] = { + "main_llm": "主问答 LLM", + "hyde_llm": "HyDE 查询增强", + "embedding": "Embedding", + "reranker": "Reranker", +} + + +@router.get("/stats") +async def get_stats(): + """Return document statistics (cached for 10 s).""" + global _stats_cache, _stats_cache_time + now = time.time() + if _stats_cache and (now - _stats_cache_time) < _STATS_TTL_SECONDS: + return _stats_cache + + documents = get_document_query_service().list_documents() + indexed = sum(1 for d in documents if d.status.value == "indexed") + failed = sum(1 for d in documents if d.status.value == "failed") + _stats_cache = { + "documents_total": len(documents), + "documents_indexed": indexed, + "documents_failed": failed, + "chunks_total": sum(d.chunk_count for d in documents), + } + _stats_cache_time = now + return _stats_cache + + +@router.get("/config") +async def get_config(): + """Return system configuration.""" + return { + "embedding_model": settings.embedding_model, + "embedding_dim": settings.embedding_dim, + "embedding_base_url": settings.embedding_base_url, + "milvus_collection": settings.milvus_collection, + "parser_backend": settings.parser_backend, + "chunk_backend": settings.chunk_backend, + "artifact_prefix": settings.document_parse_artifact_prefix, + "parser_failure_mode": settings.parser_failure_mode, + "llm_provider": settings.llm_provider, + "llm_model": settings.llm_model, + "document_metadata_path": settings.document_metadata_path, + } + + +@router.get("/milvus/health") +async def milvus_health(): + """Return Milvus health (kept for backwards compat).""" + return get_vector_index().health() + + +@router.get("/health") +async def get_health(): + """Return aggregate health of all backend services.""" + # --- Milvus --- + try: + milvus_info = get_vector_index().health() + milvus_status = "ok" if milvus_info.get("connected") else "error" + except Exception as exc: # noqa: BLE001 + milvus_info = {} + milvus_status = "error" + milvus_info["error"] = str(exc) + + # --- MinIO --- + try: + minio_connected = get_binary_store().client.connected + minio_status = "ok" if minio_connected else "error" + except Exception: # noqa: BLE001 + minio_status = "error" + minio_connected = False + + # --- BM25 --- + bm25 = get_bm25_retriever() + + # --- Sessions --- + try: + session_count = len(get_conversation_store().list_sessions()) + except Exception: # noqa: BLE001 + session_count = 0 + + return { + "milvus": {"status": milvus_status, **milvus_info}, + "minio": {"status": minio_status, "connected": minio_connected}, + "bm25": {"available": bm25 is not None}, + "reranker": { + "enabled": settings.reranker_enabled, + "model": settings.reranker_model if settings.reranker_enabled else None, + }, + "sessions": { + "active": session_count, + "max": settings.session_max_sessions, + }, + } + + +def _resolve_role_provider_model(role: str) -> tuple[str, str]: + """Return the (provider, model) pair currently configured for one AI model role. + + For "hyde_llm" this mirrors the exact fallback logic already used in + hyde_expander.py (settings.hyde_llm_provider or settings.llm_provider, same + for model) so tracker lookups here always match what TrackedLLMClient + recorded when HyDE actually ran. + """ + if role == "main_llm": + return settings.llm_provider, settings.llm_model + if role == "hyde_llm": + return ( + settings.hyde_llm_provider or settings.llm_provider, + settings.hyde_llm_model or settings.llm_model, + ) + if role == "embedding": + return "embedding", settings.embedding_model + if role == "reranker": + return "reranker", settings.reranker_model + raise ValueError(f"unknown model role: {role}") # pragma: no cover - internal roles are fixed + + +def _build_model_status(role: str) -> dict[str, Any]: + """Build one /status/models row for the given role from tracker data + live settings.""" + provider, model = _resolve_role_provider_model(role) + entry = get_model_usage_tracker().get(provider, model) + + main_provider, main_model = _resolve_role_provider_model("main_llm") + shares_usage_with = ( + "main_llm" if role != "main_llm" and (provider, model) == (main_provider, main_model) else None + ) + + enabled = True + status = entry.status if entry else "never_called" + if role == "reranker": + enabled = settings.reranker_enabled + if not enabled: + # Config always wins: report "disabled" even if the reranker was + # enabled and called successfully earlier in this process's life. + status = "disabled" + + return { + "role": role, + "role_label": _MODEL_ROLES[role], + "provider": provider, + "model": model, + "enabled": enabled, + "status": status, + "total_tokens": entry.total_tokens if entry else 0, + "call_count_ok": entry.call_count_ok if entry else 0, + "call_count_error": entry.call_count_error if entry else 0, + "last_called_at": entry.last_called_at.isoformat() if entry and entry.last_called_at else None, + "last_latency_ms": entry.last_latency_ms if entry else None, + "last_error": entry.last_error if entry else None, + "shares_usage_with": shares_usage_with, + } + + +@router.get("/models") +async def get_model_statuses(): + """Return connection status + cumulative token usage for all 4 tracked AI model roles. + + Passive: reads tracker state + settings only, makes no outbound network calls. + """ + return {"models": [_build_model_status(role) for role in _MODEL_ROLES]} + + +async def _ping_main_or_hyde(role: str) -> None: + """Send one minimal chat completion to the LLM configured for `role`.""" + provider, model = _resolve_role_provider_model(role) + client = get_llm_client(provider=provider, model=model) + await asyncio.to_thread(client.chat, [{"role": "user", "content": "ping"}], max_tokens=1) + + +async def _ping_embedding() -> None: + """Send one minimal embedding request.""" + await asyncio.to_thread(get_embedding_provider().embed_query, "ping") + + +async def _ping_reranker() -> None: + """Send one minimal rerank request, only when the reranker is enabled.""" + reranker = get_reranker() + if reranker is None: + return + # Minimal single-chunk probe — real content doesn't matter, only round-trip success. + placeholder = RetrievedChunk(chunk_id="ping", doc_id="ping", doc_title="ping", text="ping", score=0.0) + await asyncio.to_thread(reranker.rerank, "ping", [placeholder], 1) + + +@router.post("/models/ping") +async def ping_model_connections(): + """Actively test each configured model with a minimal request, then return fresh statuses. + + Each ping is isolated with return_exceptions=True so one model timing out + or erroring does not prevent the other three from completing and being + reported. Failures are still visible afterwards via _build_model_status() + because the underlying clients record their own outcome into the tracker. + """ + tasks = [ + _ping_main_or_hyde("main_llm"), + _ping_main_or_hyde("hyde_llm"), + _ping_embedding(), + _ping_reranker(), + ] + await asyncio.gather(*tasks, return_exceptions=True) + return {"models": [_build_model_status(role) for role in _MODEL_ROLES]} +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/test_status_models_routes.py -v` +Expected: PASS (4 tests) + +Then run the existing status-adjacent tests to confirm no regressions (auth routes still import `app.api.main` the same way): + +Run: `uv run pytest tests/test_auth_routes.py backend/tests/observability/ tests/test_status_models_routes.py -v` +Expected: PASS (all) + +- [ ] **Step 5: Commit** + +```bash +git add backend/app/api/routes/status.py tests/test_status_models_routes.py +git commit -m "feat: add GET/POST /status/models routes for AI model connection status" +``` + +--- + +### Task 7: Frontend — i18n keys + `.status.error` CSS fix + +**Files:** +- Modify: `frontend/src/locales/en.ts:2-134` (interface) and `:349-394` (value object) +- Modify: `frontend/src/locales/zh.ts:90-135` (value object) +- Modify: `frontend/src/styles/globals.css:133-140` + +**Interfaces:** +- Produces: new `Translations.status.*` keys (`cardModels`, `testConnectionBtn`, `testingBtn`, `roleMainLlm`, `roleHydeLlm`, `roleEmbedding`, `roleReranker`, `modelStatusNeverCalled`, `modelStatusDisabled`, `sharesUsageWithMain`, `lastCalledNever`) consumed by Task 9. New CSS rule `.status.error` (this repo's `StatusIcon`/badge components already emit `className="status error"` for error states, e.g. Milvus health, but only `.status.ok/.warn/.risk/.info` exist today — `.error` silently renders uncolored; this new card reuses that exact badge pattern, so the gap is fixed here rather than shipped a second time). + +- [ ] **Step 1: Add the type declarations (this "fails" the TypeScript build until Step 3 supplies matching values)** + +In `frontend/src/locales/en.ts`, in the `export interface Translations` block, inside the `status: { ... }` section, insert after the existing `totalChunks: string;` line (line 133): + +```typescript + totalChunks: string; + cardModels: string; + testConnectionBtn: string; + testingBtn: string; + roleMainLlm: string; + roleHydeLlm: string; + roleEmbedding: string; + roleReranker: string; + modelStatusNeverCalled: string; + modelStatusDisabled: string; + sharesUsageWithMain: string; + lastCalledNever: string; + }; +``` + +(This replaces the original `totalChunks: string;\n };` two-line ending of the `status` type block with the expanded version above.) + +- [ ] **Step 2: Run the frontend build to verify it fails** + +Run: `npm --prefix frontend run build` +Expected: FAIL — TypeScript error `Property 'cardModels' is missing in type '{ ... }' but required in type 'Translations'` (or similar), pointing at `zh.ts`'s `export const zh: Translations = {...}` object, because `zh.ts` does not yet have these keys. + +- [ ] **Step 3: Supply matching values in both locale files** + +In `frontend/src/locales/en.ts`, in the `export const en: Translations = {...}` value object, inside its `status: { ... }` section, insert after the existing `totalChunks: 'Total vector chunks',` line (line 393): + +```typescript + totalChunks: 'Total vector chunks', + cardModels: 'AI Models', + testConnectionBtn: 'Test connection', + testingBtn: 'Testing…', + roleMainLlm: 'Main answer LLM', + roleHydeLlm: 'HyDE query expansion', + roleEmbedding: 'Embedding', + roleReranker: 'Reranker', + modelStatusNeverCalled: 'Not called yet', + modelStatusDisabled: 'Disabled', + sharesUsageWithMain: 'Shares usage with main LLM', + lastCalledNever: 'Never', + }, +``` + +In `frontend/src/locales/zh.ts`, in the `export const zh: Translations = {...}` value object, inside its `status: { ... }` section, insert after the existing `totalChunks: '向量分块总数',` line (line 134): + +```typescript + totalChunks: '向量分块总数', + cardModels: 'AI 模型', + testConnectionBtn: '测试连接', + testingBtn: '测试中…', + roleMainLlm: '主问答 LLM', + roleHydeLlm: 'HyDE 查询增强', + roleEmbedding: 'Embedding', + roleReranker: 'Reranker', + modelStatusNeverCalled: '尚未调用', + modelStatusDisabled: '已禁用', + sharesUsageWithMain: '与主 LLM 共用统计', + lastCalledNever: '从未', + }, +``` + +- [ ] **Step 4: Run the frontend build to verify it passes** + +Run: `npm --prefix frontend run build` +Expected: PASS (TypeScript compiles; the two locale files now satisfy `Translations`) + +- [ ] **Step 5: Add the missing `.status.error` CSS rule** + +In `frontend/src/styles/globals.css`, after the existing block (lines 137-138): + +```css +.status.risk { color: var(--danger); background: var(--danger-bg); } +.status.risk::before { background: var(--danger); } +``` + +insert: + +```css +.status.error { color: var(--danger); background: var(--danger-bg); } +.status.error::before { background: var(--danger); } +``` + +- [ ] **Step 6: Run lint + build once more to confirm nothing broke** + +Run: `npm --prefix frontend run lint` +Expected: PASS (no new errors) + +Run: `npm --prefix frontend run build` +Expected: PASS + +- [ ] **Step 7: Commit** + +```bash +git add frontend/src/locales/en.ts frontend/src/locales/zh.ts frontend/src/styles/globals.css +git commit -m "feat: add i18n keys for AI Models card and fix missing .status.error CSS" +``` + +--- + +### Task 8: Frontend — API client layer (`api/index.ts`, `api/status.ts`) + +**Files:** +- Modify: `frontend/src/api/index.ts:279-313` (add types before the closing `export { API_BASE_URL };`) +- Modify: `frontend/src/api/status.ts` (full file, 16 lines) + +**Interfaces:** +- Consumes: `fetchAPI` (`frontend/src/api/index.ts`, existing), the `GET /status/models` / `POST /status/models/ping` response shape (Task 6). +- Produces: `ModelUsageEntry` / `ModelUsageResponse` types; `getModelUsage(): Promise`; `pingModelConnections(): Promise` — consumed by Task 9. + +- [ ] **Step 1: Add types to `frontend/src/api/index.ts`** + +Insert before the final `export { API_BASE_URL };` line: + +```typescript +export type ModelRole = 'main_llm' | 'hyde_llm' | 'embedding' | 'reranker'; +export type ModelStatus = 'ok' | 'error' | 'never_called' | 'disabled'; + +export interface ModelUsageEntry { + role: ModelRole; + role_label: string; + provider: string; + model: string; + enabled: boolean; + status: ModelStatus; + total_tokens: number; + call_count_ok: number; + call_count_error: number; + last_called_at: string | null; + last_latency_ms: number | null; + last_error: string | null; + shares_usage_with: ModelRole | null; +} + +export interface ModelUsageResponse { + models: ModelUsageEntry[]; +} +``` + +- [ ] **Step 2: Add functions to `frontend/src/api/status.ts`** + +Replace the full contents of `frontend/src/api/status.ts`: + +```typescript +import { fetchAPI, type ModelUsageResponse, type SystemConfig, type SystemHealth, type SystemStats } from './index'; + +export async function getSystemStats(): Promise { + return fetchAPI('/status/stats'); +} + +export async function getSystemConfig(): Promise { + return fetchAPI('/status/config'); +} + +export async function getSystemHealth(): Promise { + return fetchAPI('/status/health'); +} + +/** Passive read: current connection status + cumulative token usage for all 4 AI model roles. */ +export async function getModelUsage(): Promise { + return fetchAPI('/status/models'); +} + +/** Active check: sends one minimal request to each enabled model, then returns fresh statuses. */ +export async function pingModelConnections(): Promise { + return fetchAPI('/status/models/ping', { method: 'POST' }); +} + +export type { ModelUsageResponse, SystemConfig, SystemHealth, SystemStats }; +``` + +- [ ] **Step 3: Run the frontend build to verify it compiles** + +Run: `npm --prefix frontend run build` +Expected: PASS (new exports compile; nothing consumes them yet, which is fine — Task 9 wires them in) + +- [ ] **Step 4: Commit** + +```bash +git add frontend/src/api/index.ts frontend/src/api/status.ts +git commit -m "feat: add ModelUsageEntry types and status API client functions" +``` + +--- + +### Task 9: Frontend — "AI Models" card on `StatusPage.tsx` + +**Files:** +- Modify: `frontend/src/pages/Status/StatusPage.tsx` (full file, 366 lines) + +**Interfaces:** +- Consumes: `getModelUsage()`, `pingModelConnections()`, `ModelUsageEntry`, `ModelUsageResponse` (Task 8); `t.status.cardModels` etc. (Task 7); existing `StatusIcon` component (defined in this same file) and `.service-row`/`.status.*` CSS classes. + +- [ ] **Step 1: Add the import and new state** + +At the top of `frontend/src/pages/Status/StatusPage.tsx`, change: + +```typescript +import { useState, useEffect } from 'react'; +import { Topbar } from '../../components/layout/Topbar'; +import { Search, Upload, Download, RefreshCw, CheckCircle, XCircle, AlertTriangle, Info } from 'lucide-react'; +import { UploadModal } from '../Docs/UploadModal'; +import { useLanguage } from '../../contexts/LanguageContext'; +``` + +to: + +```typescript +import { useState, useEffect } from 'react'; +import { Topbar } from '../../components/layout/Topbar'; +import { Search, Upload, Download, RefreshCw, CheckCircle, XCircle, AlertTriangle, Info } from 'lucide-react'; +import { UploadModal } from '../Docs/UploadModal'; +import { useLanguage } from '../../contexts/LanguageContext'; +import { getModelUsage, pingModelConnections } from '../../api/status'; +import type { ModelUsageEntry } from '../../api/index'; +``` + +(Confirmed against the existing sibling import in `frontend/src/pages/RagChat/RagChatPage.tsx:8` — `import type { SSEMessage } from '../../api/index';` — this file uses the explicit `/index` suffix, not the bare `'../../api'` directory import, so the line above matches established convention exactly.) + +Inside `export function StatusPage()`, after the existing `const [lastRefresh, setLastRefresh] = useState(null);` line, add: + +```typescript + const [modelUsage, setModelUsage] = useState(null); + const [pinging, setPinging] = useState(false); +``` + +- [ ] **Step 2: Fetch model usage alongside the other 3 endpoints** + +Change the existing `useEffect` fetch block from: + +```typescript + // Fetch all three endpoints in parallel + Promise.allSettled([ + fetch('/api/v1/status/stats', { headers: authHeader() }).then(r => r.json()), + fetch('/api/v1/status/health', { headers: authHeader() }).then(r => r.json()), + fetch('/api/v1/status/config', { headers: authHeader() }).then(r => r.json()), + ]).then(([statsRes, healthRes, configRes]) => { + if (statsRes.status === 'fulfilled') setStats(statsRes.value); + else setStats({ documents_total: 0, documents_indexed: 0, documents_failed: 0, chunks_total: 0 }); + + if (healthRes.status === 'fulfilled') setHealth(healthRes.value); + if (configRes.status === 'fulfilled') setConfig(configRes.value); + + setLoading(false); + setHealthLoading(false); + setLastRefresh(new Date()); + }); +``` + +to: + +```typescript + // Fetch all endpoints in parallel. The first three use raw fetch() (legacy + // pattern already established in this file); model usage uses the typed + // fetchAPI-based client from api/status.ts — new code should prefer that. + Promise.allSettled([ + fetch('/api/v1/status/stats', { headers: authHeader() }).then(r => r.json()), + fetch('/api/v1/status/health', { headers: authHeader() }).then(r => r.json()), + fetch('/api/v1/status/config', { headers: authHeader() }).then(r => r.json()), + getModelUsage(), + ]).then(([statsRes, healthRes, configRes, modelsRes]) => { + if (statsRes.status === 'fulfilled') setStats(statsRes.value); + else setStats({ documents_total: 0, documents_indexed: 0, documents_failed: 0, chunks_total: 0 }); + + if (healthRes.status === 'fulfilled') setHealth(healthRes.value); + if (configRes.status === 'fulfilled') setConfig(configRes.value); + if (modelsRes.status === 'fulfilled') setModelUsage(modelsRes.value.models); + else setModelUsage(null); + + setLoading(false); + setHealthLoading(false); + setLastRefresh(new Date()); + }); +``` + +- [ ] **Step 3: Add the "test connection" handler** + +After the existing `function handleExport() { ... }` function, add: + +```typescript + async function handleTestConnections() { + setPinging(true); + try { + const res = await pingModelConnections(); + setModelUsage(res.models); + } catch { + // Leave modelUsage as-is; the card below already shows a muted + // "never_called"/error state per row when data can't be refreshed. + } finally { + setPinging(false); + } + } + + function modelBadgeStatus(status: ModelUsageEntry['status']): 'ok' | 'error' | 'warn' | 'info' { + if (status === 'ok') return 'ok'; + if (status === 'error') return 'error'; + if (status === 'disabled') return 'info'; + return 'info'; // never_called + } + + function modelStatusLabel(entry: ModelUsageEntry): string { + if (entry.status === 'never_called') return t.status.modelStatusNeverCalled; + if (entry.status === 'disabled') return t.status.modelStatusDisabled; + return entry.status === 'ok' ? t.status.badgeOnline : t.status.badgeError; + } + + /** Small relative-ish hint shown next to provider/model — "Never" or a local time string. */ + function modelLastCalledLabel(entry: ModelUsageEntry): string { + if (!entry.last_called_at) return t.status.lastCalledNever; + return new Date(entry.last_called_at).toLocaleTimeString(undefined, { hour: '2-digit', minute: '2-digit' }); + } +``` + +- [ ] **Step 4: Render the new card** + +Immediately after the closing `` of the existing "System health" card (the `
` block containing `{t.status.cardHealth}` — locate it by searching for `cardHealth` in this file) and before the "System config (collapsible)" card's opening `
`, insert: + +```typescript + {/* AI Models — connection status + cumulative token usage */} +
+
+ {t.status.cardModels} + +
+ + {!modelUsage ? ( +
+ {[1, 2, 3, 4].map(i => ( +
+ ))} +
+ ) : ( + modelUsage.map(entry => { + const roleLabel = entry.role === 'main_llm' ? t.status.roleMainLlm + : entry.role === 'hyde_llm' ? t.status.roleHydeLlm + : entry.role === 'embedding' ? t.status.roleEmbedding + : t.status.roleReranker; + return ( +
+ + {roleLabel} + + {entry.provider}/{entry.model} + {entry.shares_usage_with && ` · ${t.status.sharesUsageWithMain}`} + {` · ${modelLastCalledLabel(entry)}`} + + + {entry.total_tokens > 0 || entry.status === 'ok' || entry.status === 'error' + ? entry.total_tokens.toLocaleString() + : '—'} + + + {modelStatusLabel(entry)} + +
+ ); + }) + )} +
+``` + +- [ ] **Step 5: Remove the now-duplicate reranker line from the "Runtime" card** + +In the "Sessions & reranker quick facts" card (search for `cardRuntime` in this file), remove this block: + +```typescript +
+ {t.status.labelReranker} + + {health.reranker.enabled ? (health.reranker.model ?? t.status.serviceEnabled) : t.status.serviceDisabled} + +
+``` + +(leaving `labelActiveSessions`, `labelSessionCapacity`, and `labelBM25` rows in place — only the reranker row is removed, since the new "AI Models" card now shows it with richer detail.) + +- [ ] **Step 6: Run lint + build + manual check** + +Run: `npm --prefix frontend run lint` +Expected: PASS + +Run: `npm --prefix frontend run build` +Expected: PASS + +Run: `npm --prefix frontend run dev` (or `./dev.sh start frontend --foreground` per `AGENTS.md`), open the Status page in a browser, and manually verify: +- The new "AI Models" card renders 4 rows (main LLM, HyDE, Embedding, Reranker) without console errors. +- Clicking "Test connection" shows "Testing…" then updates the rows (requires a reachable backend + gateway; if the gateway is unreachable, at least confirm the button doesn't crash the page and an error-status row renders in red, validating the CSS fix from Task 7). +- The "Runtime" card no longer shows a separate Reranker line. + +- [ ] **Step 7: Commit** + +```bash +git add frontend/src/pages/Status/StatusPage.tsx +git commit -m "feat: add AI Models card to Status page with connection test button" +``` + +--- + +## Self-Review + +**1. Spec coverage:** +- Tracker keyed by provider:model, never-raises → Task 1. ✅ +- TrackedLLMClient wrapping every call site via the factory → Tasks 2-3. ✅ +- Embedding usage capture, reranker call-only tracking → Tasks 4-5. ✅ +- `GET /status/models` (passive) + `POST /status/models/ping` (active, parallel, partial-failure-safe) → Task 6. ✅ +- Frontend card, "shares usage with main LLM" note, removed duplicate Runtime reranker line, i18n, desktop-only → Tasks 7-9. ✅ +- Known limitation (streaming tokens not counted) → encoded directly in `TrackedLLMClient.stream_chat()`'s docstring/comment (Task 2). ✅ +- `.status.error` CSS gap → Task 7, Step 5. ✅ +- No new DB table, no cost estimation, no per-session breakdown → correctly absent from every task. ✅ + +**2. Placeholder scan:** No "TBD"/"TODO"/"add appropriate handling" phrases in any step; every code block is complete, runnable code; every test has real assertions. Task 9's import path was independently verified against `frontend/src/pages/RagChat/RagChatPage.tsx:8` (an existing sibling file) rather than left as a "verify this" note for the executor. + +**3. Type consistency:** +- `ModelUsageEntry.status` values (`never_called|ok|error`) plus the route's `disabled` override are consistently referenced as the 4-value `ModelStatus` union across Task 1 (Python), Task 6 (route), and Task 8 (TypeScript type) — checked field-by-field. +- `ModelUsageTracker.record()`'s keyword-only signature (`provider, model, success, usage=None, latency_ms=None, error=None`) is called identically in Tasks 2, 4, 5, and matches Task 1's definition exactly (checked argument names one by one). +- `TrackedLLMClient` is used, not subclassed, everywhere (Task 3's `llm_factory.py` return-type annotations updated to the union `BaseLLMClient | TrackedLLMClient` consistently in `create()`, `get_cached()`, and `get_llm_client()`). +- Frontend `ModelUsageEntry`/`ModelUsageResponse` (Task 8) field names match the backend route's dict keys (Task 6) one-to-one: `role, role_label, provider, model, enabled, status, total_tokens, call_count_ok, call_count_error, last_called_at, last_latency_ms, last_error, shares_usage_with`. +- i18n keys referenced in Task 9's JSX (`t.status.cardModels`, `testConnectionBtn`, `testingBtn`, `roleMainLlm`, `roleHydeLlm`, `roleEmbedding`, `roleReranker`, `modelStatusNeverCalled`, `modelStatusDisabled`, `sharesUsageWithMain`, `lastCalledNever`) all match keys added in Task 7 exactly (checked name-by-name) and all are actually rendered — the earlier draft had an unused `totalTokensLabel` key and omitted the spec's "last-called relative-time hint"; both are fixed: the unused key was removed from Task 7, and Task 9 now renders `modelLastCalledLabel()` (using `lastCalledNever`) inline next to provider/model.