AIRegulation-DocAnalysis/docs/superpowers/specs/2026-07-23-status-model-usage-hardening-design.md

System Status — AI Models Panel Hardening Design

Date: 2026-07-23 Scope: Close three gaps left open by the already-shipped "AI Models" card on the System Status page (docs/superpowers/specs/2026-07-02-status-llm-model-usage-design.md): streaming calls don't report token usage, the Cross-Encoder reranker is still disabled, and usage counters reset on every backend restart. Relationship to existing roadmap: This is a direct continuation of the 2026-07-02 feature, not a new module. It also closes two long-standing "Quick Win" items from AI_Agent_优化分析报告_2026-06-18.md (reranker enablement, and — partially — observability of RAG quality). It does not attempt full Langfuse/Ragas tracing (P0-A in that roadmap); that remains a separate, larger effort.

---

Goals

1. A1 — Streaming token capture. stream_chat() calls (the default interaction mode for the main RAG chat UI) currently report call success/latency but not token usage — an explicitly documented gap in the 2026-07-02 design. Close it using the OpenAI-compatible stream_options: {include_usage: true} mechanism, so streaming and non-streaming calls are accounted for consistently.
2. A2 — Enable the reranker. reranker_enabled has been False by default since before the first internal analysis report (2026-06-11); both that report and the 2026-06-18 follow-up flag it as the single highest-ROI, lowest-risk unfinished item (+15–25% retrieval precision, typically a one-line config change).
3. A3 — Durable usage counters. ModelUsageTracker is in-memory only; counts reset on every restart/redeploy. Persist them so the Status page reflects cumulative usage across the process lifetime, not just since the last restart.

Non-Goals

• Cost/spend estimation in currency (still no reliable pricing for the internal gateway).
• Per-session/per-user token attribution.
• Historical time-series / usage-over-time charts (explicitly deferred by user decision during brainstorming — this iteration persists current cumulative counters only, not a time-series log).
• Full Langfuse/Ragas distributed tracing and faithfulness scoring (P0-A, separate future effort).

---

A1 — Streaming Token Capture

Current behavior (confirmed by reading the code)

DeepSeekClient.stream_chat() and both QwenClient.stream_chat() / QwenVLClient.stream_chat() (backend/app/services/llm/deepseek_client.py, backend/app/services/llm/qwen_client.py) parse each SSE data: {...} line, and today explicitly skip any chunk whose choices array is empty:

choices = data.get("choices", [])
if not choices:
    continue   # <- a trailing usage-only chunk is silently dropped here today
delta = choices[0].get("delta", {})
content = delta.get("content", "")

TrackedLLMClient.stream_chat() (backend/app/services/llm/tracked_client.py) wraps this with a plain for chunk in self._inner.stream_chat(...): yield chunk, then records call success/latency only — by design, since "none of the current provider stream_chat() implementations parse a trailing usage chunk."

Change

1. Add "stream_options": {"include_usage": True} to the request payload built in each of the three stream_chat() implementations. This is the standard OpenAI-compatible mechanism: the gateway appends one final chunk with "choices": [] and a populated "usage" object after the normal content chunks.
2. In each generator, when a parsed chunk has empty choices and a non-empty usage field, capture it into a local variable (function-local — safe even though the underlying client instance is a shared/cached singleton, because each call to stream_chat() creates its own generator frame). At the end of the generator, return that captured usage dict instead of falling off the end with an implicit None. This is accessible to a manual consumer via StopIteration.value.
3. Per-chunk content yields are unchanged — this keeps the change backward compatible for all seven existing call sites (api/routes/rag.py, compliance.py, agent.py, application/perception/services.py, infrastructure/llm/openai_compatible_answer_generator.py, services/agent/qa_agent.py) that just do for chunk in stream_chat(...): ... and will continue to work untouched, silently ignoring the new return value.
4. TrackedLLMClient.stream_chat() is the only call site that needs the return value. Replace its plain for loop with a manually-driven loop (next() in a try/except StopIteration) so it can capture StopIteration.value and pass it into the same existing self._tracker.record(...) call in its finally block — no new tracker call, no double-counting of call_count_ok/call_count_error.

Known limitation carried forward

If the gateway does not honor stream_options.include_usage (some OpenAI-compatible proxies ignore unknown fields silently rather than erroring), streaming usage will simply remain absent, same as today — this is a graceful no-op, not a new failure mode.

---

A2 — Enable the Reranker

Current behavior (confirmed)

.env has RERANKER_ENABLED=false. OpenAICompatibleReranker (backend/app/infrastructure/vectorstore/cross_encoder_reranker.py) already:

• Tries TEI format (POST /rerank), falls back to Cohere format (POST /v1/rerank) on 400/404.
• On any failure, logs a warning, records the failure into ModelUsageTracker (provider="reranker"), and falls back to the original unranked order rather than raising — retrieval keeps working even if the reranker is broken.

Change

Flip RERANKER_ENABLED=true in root .env. No code change. rag_retrieval_top_k=20 / rag_top_k=5 are already set to reasonable pre/post-rerank values (backend/app/config/settings.py:124-125).

Verification

Use the already-shipped POST /status/models/ping endpoint to actively confirm the gateway's rerank endpoint responds before considering this done. If it errors, the Status page's "AI Models" card will show the reranker row as error (existing behavior, not new) — revert the flag in that case rather than leaving retrieval silently degraded to "reranker enabled but always failing over."

---

A3 — Durable Usage Counters

Current behavior (confirmed)

ModelUsageTracker (backend/app/shared/model_usage_tracker.py) holds all state in an in-memory dict guarded by a threading.Lock. Nothing writes it to disk; a restart or redeploy zeroes every counter.

Design decision (confirmed with user during brainstorming)

Persist current cumulative counters only — one row per provider:model, no historical/time-series log. This is the smaller, already-shaped slice of what the 2026-07-02 spec deferred; a time-series log can be layered on top later if trend charts are ever requested, without reworking this table.

Data model

New table, created the same way every other Postgres store in this codebase creates its table — a CREATE TABLE IF NOT EXISTS string executed on first use, no migration framework (matches postgres_event_store.py, postgres_document_repository.py, postgres_document_processing_store.py, user_store.py, compliance/repository.py — all follow this idiom):

CREATE TABLE IF NOT EXISTS model_usage_stats (
    provider          VARCHAR(64)  NOT NULL,
    model             VARCHAR(128) NOT NULL,
    total_tokens      BIGINT       NOT NULL DEFAULT 0,
    prompt_tokens     BIGINT       NOT NULL DEFAULT 0,
    completion_tokens BIGINT       NOT NULL DEFAULT 0,
    call_count_ok     BIGINT       NOT NULL DEFAULT 0,
    call_count_error  BIGINT       NOT NULL DEFAULT 0,
    last_called_at    TIMESTAMPTZ,
    last_latency_ms   INTEGER,
    last_error        TEXT,
    updated_at        TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
    PRIMARY KEY (provider, model)
);

Gating — reuse the existing backend toggle, don't add a new one

Gate persistence behind the existing settings.document_repository_backend == "postgres" flag (the same one documents/compliance already key off in backend/app/shared/bootstrap.py), rather than introducing a new setting. When it is "json" (today's default), ModelUsageTracker behaves exactly as it does today — purely in-memory, zero new hard requirement on a running Postgres for local/dev use.

Write strategy: periodic snapshot flush, not per-call write-through

Single-worker assumption: this design assumes a single backend worker process. In a multi-worker deployment (e.g., multiple Uvicorn workers or replicas), each worker holds its own in-memory ModelUsageTracker and its own periodic flush will overwrite the same (provider, model) row with only that worker's partial counts (last-writer-wins semantics), so persisted totals would under-count versus true cross-worker totals — a pre-existing limitation of ModelUsageTracker being per-process, now also reflected in what gets persisted.

Rejected: writing to Postgres synchronously inside record() on every single LLM/embedding/reranker call — this would add blocking DB I/O to the hot path of every chat/RAG/compliance request, contradicting the tracker's own documented principle that tracking "must never disrupt a real user-facing call."

Chosen approach:

• On startup (in the existing lifespan() hook in backend/app/api/main.py, alongside the existing preload_runtime_dependencies() call): if the postgres backend is active, load existing model_usage_stats rows and seed ModelUsageTracker's in-memory dict, so counts continue cumulatively instead of restarting at zero.
• Every 60 seconds, a background asyncio task (started at the same point, cancelled in the existing shutdown/cleanup_runtime_dependencies() path) snapshots the tracker (tracker.snapshot(), already exists) and UPSERTs each entry (INSERT ... ON CONFLICT (provider, model) DO UPDATE) — overwriting with the current cumulative value, not incrementing, so a missed cycle is never double-counted.
• Best-effort flush on shutdown as a bonus on top of the periodic flush (not the primary durability mechanism — a SIGKILL/OOM crash will not trigger it, which is an acceptable, explicitly-noted gap for an observability feature: worst case, up to 60s of counters are lost, not corrupted).

---

Error Handling

• A1: if a client's stream_chat() never emits a trailing usage chunk (gateway doesn't support stream_options), the generator simply returns None; TrackedLLMClient already treats "no usage" as a no-op for the token fields (existing record() behavior — usage or {}).
• A2: unchanged — already-shipped graceful fallback and error surfacing.
• A3: the flush task wraps each cycle in try/except Exception: logger.warning(...) — a transient Postgres blip must not crash the flush loop or the app; it simply retries on the next 60s tick. Startup load failure (e.g., Postgres unreachable at boot) logs a warning and leaves the tracker empty, exactly as it behaves today with no persistence at all — it does not block app startup.

Testing

Mirrors existing conventions (backend/tests/observability/, backend/tests/perception/test_postgres_event_store.py for the mocked-psycopg2 pattern — no real database needed):

• Extend backend/tests/observability/test_tracked_client.py: streaming usage now flows into the same record() call (assert the returned StopIteration.value path is wired correctly).
• New tests for DeepSeekClient.stream_chat() / QwenClient.stream_chat() / QwenVLClient.stream_chat(): trailing usage chunk is parsed and returned; ordinary content chunks are unaffected; a stream with no usage chunk still returns None without error.
• New backend/tests/observability/test_model_usage_persistence.py: mocked psycopg2 (same pattern as test_postgres_event_store.py) — verifies startup load seeds the tracker, the flush cycle upserts a snapshot, and everything is a no-op when document_repository_backend != "postgres".
• A2 needs no new test — it is a configuration change exercised by existing reranker tests and manual /status/models/ping verification.

Out of Scope (deferred to future iterations)

• Time-series/historical usage log and trend charts (explicit user decision this iteration — durable counters only).
• Cost/spend estimation in currency.
• Per-session/per-user attribution.
• Full Langfuse/Ragas tracing and LLM-as-judge faithfulness scoring (P0-A, tracked separately).