wangweiandCopilot
0907470a2d
Fix 1 (bootstrap.py): wrap store.flush() in asyncio.to_thread() inside the periodic _flush_loop() to avoid blocking the async event loop every 60s. Synchronous signatures of _start/_stop_model_usage_persistence() and the one-time seed/shutdown flushes are left unchanged per review scope. Fix 2 (test_stream_chat_usage_capture.py): add the two-line stream_options assertion to test_qwen_vl_stream_chat_returns_usage_from_trailing_chunk, matching the identical check already present in the DeepSeek and Qwen tests. Fix 3 (design doc): note the single-worker assumption in A3's write strategy section — multi-worker deployments get last-writer-wins per-row semantics. Tests: 69 passed, 0 failed (python -m pytest backend/tests -q) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
140 lines
12 KiB
Markdown
140 lines
12 KiB
Markdown
# 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).
|