diff --git a/backend/app/shared/bootstrap.py b/backend/app/shared/bootstrap.py index 36c07c4..74836fd 100644 --- a/backend/app/shared/bootstrap.py +++ b/backend/app/shared/bootstrap.py @@ -464,7 +464,7 @@ def _start_model_usage_persistence() -> None: while True: await asyncio.sleep(60) try: - store.flush(tracker.snapshot()) + await asyncio.to_thread(store.flush, tracker.snapshot()) except Exception as exc: # noqa: BLE001 - one bad cycle must not kill the loop logger.warning("Failed to flush model usage stats: {}", exc) diff --git a/backend/tests/observability/test_stream_chat_usage_capture.py b/backend/tests/observability/test_stream_chat_usage_capture.py index 8318f70..1ea1a83 100644 --- a/backend/tests/observability/test_stream_chat_usage_capture.py +++ b/backend/tests/observability/test_stream_chat_usage_capture.py @@ -112,3 +112,5 @@ def test_qwen_vl_stream_chat_returns_usage_from_trailing_chunk(): assert chunks == ["Describing image"] assert returned_usage == usage + sent_payload = client._client.stream.call_args.kwargs["json"] + assert sent_payload["stream_options"] == {"include_usage": True} diff --git a/docs/superpowers/specs/2026-07-23-status-model-usage-hardening-design.md b/docs/superpowers/specs/2026-07-23-status-model-usage-hardening-design.md new file mode 100644 index 0000000..18b79df --- /dev/null +++ b/docs/superpowers/specs/2026-07-23-status-model-usage-hardening-design.md @@ -0,0 +1,139 @@ +# 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: + +```python +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): + +```sql +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 `UPSERT`s 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).