Author SHA1 Message Date
wangwei b2feaeddb4 update for mcp 2026-08-06 11:08:46 +08:00
wangweiandCopilot 31bbf80aeb feat: surface MCP server status in System Status page
Add per-tool in-memory call counters to the MCP module and a
GET /api/v1/status/mcp endpoint that joins them with the live tool
registry and endpoint config, then render it as a new card on the
System Status page with a one-click client-config copy button.

- app/mcp/stats.py: lock-guarded MCPStatsTracker (the mcp SDK runs sync
  tool bodies via anyio.to_thread.run_sync, so this is genuinely
  multi-threaded, unlike the async REST routes)
- app/mcp/server.py: instrument search_regulations, add get_mcp_status()
- app/config/settings.py: optional MCP_PUBLIC_URL override, required
  because the Vite proxy and reverse proxies rewrite the Host header
- StatusPage.tsx: MCP Server card, joins the existing parallel fetch

Counters are process-local by design; token usage is already persisted
by ModelUsageTracker since MCP calls route through ask().

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-08-03 11:37:22 +08:00
wangweiandCopilot 73e79a610d docs: design spec for MCP status panel
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-08-03 10:56:56 +08:00
wangweiandCopilot 49ee50c104 fix: harden MCP endpoint after code review
Critical: the MCP SDK auto-enables DNS-rebinding protection when its host
parameter is left at the 127.0.0.1 default, hard-coding a loopback-only Host
allow-list. Every remote client (the only deployment this feature targets) was
refused with HTTP 421 before auth or the tool ran. Now driven by a new
MCP_ALLOWED_HOSTS setting, with '*' as an explicit, logged opt-out.

Also bounds query/top_k to match AskRequest (top_k is amplified 4x downstream,
so an unbounded value was a resource-exhaustion vector), decodes the
Authorization header as latin-1 per the ASGI spec instead of raising a 500 on
malformed bytes, and returns WWW-Authenticate on 401 per RFC 7235.

Moves the psycopg2 import guard into backend/tests/conftest.py: duplicated
across four test modules, it only worked because of alphabetical collection
order, and any earlier-sorting package would have reintroduced a live
connection attempt against the production database.

Registers the mcp module in the authoritative backend architecture doc.

84 backend tests pass. Verified against a live server: allowed remote Host
returns a valid initialize result, unknown Host returns 421, missing token
returns 401 with WWW-Authenticate.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-29 17:11:54 +08:00
wangweiandCopilot bd3dc38d1d feat: add MCP server module exposing search_regulations tool
- New backend/app/mcp/ module: MCPServer instance with a single
  search_regulations tool backed by the existing AgentConversationService.
- MCPAuthMiddleware reuses existing JWT auth (no new auth mechanism).
- Mounted at /mcp/ in api/main.py via Streamable HTTP transport; wired the
  MCP session manager into the existing lifespan() via AsyncExitStack
  (app.mount() does not propagate nested ASGI lifespans automatically).
- Fixed a doubled /mcp/mcp path by setting streamable_http_path to "/"
  (MCPServer.streamable_http_app() defaults to registering its own /mcp route).
- Verified end-to-end with the real mcp Python client: list_tools() returns
  search_regulations, auth correctly 401s without or with an invalid token.
- 7 new tests, 76 total (up from 69), all passing.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-29 13:00:52 +08:00
wangweiandCopilot e78c8a989f docs: add MCP search_regulations implementation plan
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-29 11:03:28 +08:00
wangweiandCopilot 483689c1e8 docs: add MCP search_regulations server design spec
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-29 10:58:01 +08:00
wangweiandCopilot 6aaaff05f5 docs: add implementation plan for status model usage hardening
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-23 15:31:19 +08:00
wangweiandCopilot 0907470a2d fix: offload flush to thread pool, add QwenVL stream_options assert, document single-worker assumption
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>
2026-07-23 15:30:21 +08:00
wangweiandCopilot 29f79d7434 feat: seed and periodically persist model usage stats to Postgres
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-23 14:32:21 +08:00
wangweiandCopilot 5d132981ad feat: add PostgresModelUsageStore and ModelUsageTracker.seed()
- Add seed() method to ModelUsageTracker for bulk-loading persisted entries at startup
- Create PostgresModelUsageStore for persistence of model usage counters to Postgres
- Store only current cumulative snapshots (no historical time-series)
- Use standard CREATE TABLE IF NOT EXISTS idiom matching other Postgres stores
- Add comprehensive mocked unit tests for both components

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-23 14:18:22 +08:00
wangweiandCopilot 4f6cc4812e revert: disable Cross-Encoder reranker
Live verification via POST /status/models/ping (after confirming the
gateway itself is reachable -- embedding role succeeded, 1.5s latency)
shows http://6.86.80.4:30080/v1/rerank returns a fast, reproducible
'503 Service Unavailable' -- not a timeout/fluke. The gateway's model
catalog (19 models: deepseek-*, glm-*, kimi-*, qwen3*, text-embedding-v3/4)
contains no cross-encoder/rerank-capable model, confirming no reranker
service is deployed behind this gateway today. Leaving RERANKER_ENABLED=true
would be a permanent no-op (graceful fallback to unranked order every call)
plus a misleading permanent error badge on the Status page. Revert until a
reranker model is actually deployed on the gateway.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-23 14:09:33 +08:00
wangweiandCopilot 2547d04b9d chore: enable Cross-Encoder reranker
Sandbox verification via /status/models/ping was inconclusive: the gateway
(6.86.80.4:30080) is unreachable from this environment entirely (embedding
role failed with the identical connection-timeout pattern, which is a
feature that definitely works in real deployment) -- not evidence that
/rerank specifically is unsupported. Reranker code already has graceful
TEI/Cohere fallback + falls back to unranked order on any failure, so this
is zero-risk to retrieval even if misconfigured. Please verify via
POST /status/models/ping in an environment with real gateway access;
revert RERANKER_ENABLED to false if that shows a genuine error.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-23 14:03:43 +08:00
wangweiandCopilot 7adc050968 feat: record streaming token usage in TrackedLLMClient.stream_chat
Implement manual generator driving using next()/StopIteration to capture
the return value (trailing usage dict) from inner stream_chat() implementations,
enabling token tracking for streaming LLM calls.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-23 13:42:33 +08:00
wangweiandCopilot f2bd0deeb3 feat: capture streaming token usage in QwenClient and QwenVLClient
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-23 13:19:23 +08:00
wangweiandCopilot 81a6d54fff feat: capture streaming token usage in DeepSeekClient.stream_chat
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-23 11:24:31 +08:00
wangwei beddc2d976 Merge pull request 'main-ruqi' (#1) from main-ruqi into main
Reviewed-on: #1
2026-07-02 22:05:17 +08:00
wangwei 52e67b0e7b Add LLM token 2026-07-02 22:03:39 +08:00
wangweiandCopilot e3afb8a07a fix: normalize LLM provider key lookup and skip disabled HyDE ping (final review)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 21:24:22 +08:00
wangweiandCopilot 6a7fe48c4c fix: use dedicated error message for AI Models card load failure (Task 9 review)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 20:03:27 +08:00
wangweiandCopilot 0edbee07d5 fix: add distinct error state for AI Models card load failure (Task 9 review)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 18:01:04 +08:00
wangweiandCopilot 2ce4c8a289 feat: add AI Models card to Status page with connection test button
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 17:38:33 +08:00
wangweiandCopilot 39a51c9e83 feat: add ModelUsageEntry types and status API client functions
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 17:23:43 +08:00
wangweiandCopilot 049da2297b feat: add i18n keys for AI Models card and fix missing .status.error CSS
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 17:06:13 +08:00
wangweiandCopilot d83286edd4 fix: honor hyde_enabled toggle and record ping failures before client creation (Task 6 review)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 16:42:58 +08:00
wangweiandCopilot 169911ab46 feat: add GET/POST /status/models routes for AI model connection status
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 16:20:35 +08:00
wangweiandCopilot 66fc388bfb feat: record reranker call outcome into ModelUsageTracker
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 15:42:20 +08:00
wangwei 41096369d3 feat: record embedding call usage into ModelUsageTracker 2026-07-02 15:15:03 +08:00
wangwei 4fea159f5b feat: wrap LLM clients with TrackedLLMClient in LLMFactory 2026-07-02 15:03:12 +08:00
wangweiandCopilot d460397dda fix: add missing test comment for backend commenting standard (Task 2 review)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 14:57:55 +08:00
wangweiandCopilot 37ea27fcbe feat: add TrackedLLMClient decorator for transparent usage recording
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 14:49:06 +08:00
wangwei 74f327c85e feat: add ModelUsageTracker for per-model token/connection tracking 2026-07-02 14:41:21 +08:00
wangweiandCopilot 4b451ef97c docs: add implementation plan for System Status AI model usage tracking
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 14:13:58 +08:00
wangweiandCopilot 55ba922250 docs: add design spec for System Status AI model connection/token usage panel
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 13:39:43 +08:00
wangwei 9212747e1b update for 1. 优化 2.中英切换 2026-06-10 11:10:36 +08:00
wangwei e7963b267e fix somethings 2026-06-08 11:16:28 +08:00
wangwei 9fea9c6a53 1. Add 登陆功能
2. 调整字体大小
3. 新增部分功能
2026-06-05 18:00:31 +08:00
167 changed files with 29320 additions and 1156 deletions
+57 -2
View File
@@ -48,8 +48,23 @@ CHUNK_OVERLAP=50
MAX_FILE_SIZE_MB=100
PARSER_BACKEND=aliyun
CHUNK_BACKEND=aliyun
# 文档元数据存储后端:json(默认)或 postgres
DOCUMENT_REPOSITORY_BACKEND=json
# 文档元数据存储后端:启用 postgres 以激活合规分析历史记录(Direction B)及 Finding Chat 持久化(Direction C
DOCUMENT_REPOSITORY_BACKEND=postgres
# Set to true only when a Celery worker is actually running (./dev.sh start worker).
# Default false: processing runs in FastAPI's threadpool — no external worker needed.
USE_CELERY_WORKER=false
# ===== 法规感知爬取配置 =====
# 单次 HTTP 请求超时(秒),含正文抓取(fetch_full_text)。
PERCEPTION_CRAWL_TIMEOUT_SECONDS=120
# 每个数据源单次爬取的最大条目数。
PERCEPTION_MAX_EVENTS_PER_SOURCE=100
# 变更判定的次要闸门:段落改动字符占比达到该阈值才送 LLM 分类。
# 数字变化(如 30米->20米)或情态词变化(应当/宜/不得等)无视此阈值,始终判定为显著变更。
PERCEPTION_DIFF_MIN_CHANGE_RATIO=0.02
# 定时全量爬取的执行间隔(秒),默认 21600 = 6 小时。
# 仅当 Celery Beat 进程在运行时才生效(./dev.sh start beat),Beat 未启动则完全不会自动爬取。
PERCEPTION_CRAWL_INTERVAL_SECONDS=21600
# ===== API配置 =====
API_HOST=0.0.0.0
@@ -92,3 +107,43 @@ ALIYUN_LLM_ENHANCEMENT=true
ALIYUN_ENHANCEMENT_MODE=VLM
DOCUMENT_PARSE_ARTIFACT_PREFIX=artifacts
PARSER_FAILURE_MODE=fail
# ===== Reranker 配置 =====
RERANKER_ENABLED=false
RERANKER_BASE_URL=http://6.86.80.4:30080/v1
RERANKER_MODEL=BAAI/bge-reranker-v2-m3
RERANKER_API_KEY=sk-fVr9KmDZNC4pGDBQj0EUWz9bDmFzNxjYC9EzZpe2bVDsxtz8
RERANKER_TOP_K=5
# ===== 会话持久化 =====
SESSION_BACKEND=redis
# ===== 认证配置 =====
# 生产环境请修改为强随机密钥: python -c "import secrets; print(secrets.token_hex(32))"
AUTH_SECRET_KEY=ai-compliance-hub-jwt-secret-2026-tsystems
AUTH_ALGORITHM=HS256
AUTH_TOKEN_EXPIRE_MINUTES=480
AUTH_ENABLED=true
# ===== CORS =====
CORS_ALLOW_ORIGINS=http://localhost:5173
# ===== HyDE ???? =====
HYDE_ENABLED=true
HYDE_MAX_TOKENS=200
HYDE_LLM_PROVIDER=qwen
HYDE_LLM_MODEL=qwen3.6-flash
# ===== MCP 服务配置 =====
# MCP SDK 在传输层绑定回环地址时会自动启用 DNS 重绑定防护:Host 头不在下表内的
# 请求一律返回 HTTP 421,且发生在进入工具逻辑之前。部署在 6.86.80.9 必须显式列出
# 该地址,否则所有远程 MCP 客户端(Claude Desktop / IDE 等)100% 连不上。
# 语法:`:*` 后缀匹配任意端口;填 `*` 表示彻底关闭该防护(不推荐)。
MCP_ALLOWED_HOSTS=6.86.80.9:*,127.0.0.1:*,localhost:*,[::1]:*
# 系统状态页 MCP 卡片展示、以及"复制接入配置"按钮写入的对外访问地址。
# 留空则由后端从请求 Host 头推导;但前端经 Vite 代理(changeOrigin: true)转发后
# Host 会被改写成 API_HOST:API_PORT,推导结果是 0.0.0.0/127.0.0.1,客户端无法使用,
# 因此远程部署必须显式指定。结尾的斜杠不能省略。
MCP_PUBLIC_URL=http://6.86.80.9:8000/mcp/
+1 -1
View File
@@ -31,5 +31,5 @@ POSTGRES_PASSWORD=postgresql123456
POSTGRES_DB=compliance_db
# ===== 文档元数据后端 =====
# 改为 postgres 以启用 PG 持久化(structure_nodes + semantic_blocks 入库
# 改为 postgres 以启用合规分析历史记录(Direction B)和 Finding ChatDirection C
DOCUMENT_REPOSITORY_BACKEND=json
+85 -4
View File
@@ -50,7 +50,26 @@ DOCUMENT_METADATA_PATH=backend/data/documents.json
PARSER_BACKEND=aliyun
CHUNK_BACKEND=aliyun
# 文档元数据存储后端:json(默认,无需数据库)或 postgres(启用 PG 持久化)
# ⚠ 以下功能需要 postgres(设为 json 时功能静默降级或报 500):
# - Direction B: 合规分析历史记录 (/compliance/history/*)
# - Direction B: DOCX 报告下载
# - Direction C: Finding Chat 消息持久化
DOCUMENT_REPOSITORY_BACKEND=json
# Set to true only when a Celery worker is running (./dev.sh start worker).
# Default false: document processing runs in FastAPI's threadpool (no external worker needed).
USE_CELERY_WORKER=false
# ===== 法规感知爬取配置 =====
# 单次 HTTP 请求超时(秒),含正文抓取(fetch_full_text)。
PERCEPTION_CRAWL_TIMEOUT_SECONDS=120
# 每个数据源单次爬取的最大条目数。
PERCEPTION_MAX_EVENTS_PER_SOURCE=100
# 变更判定的次要闸门:段落改动字符占比达到该阈值才送 LLM 分类。
# 数字变化(如 30米->20米)或情态词变化(应当/宜/不得等)无视此阈值,始终判定为显著变更。
PERCEPTION_DIFF_MIN_CHANGE_RATIO=0.02
# 定时全量爬取的执行间隔(秒),默认 21600 = 6 小时。
# 仅当 Celery Beat 进程在运行时才生效(./dev.sh start beat),Beat 未启动则完全不会自动爬取。
PERCEPTION_CRAWL_INTERVAL_SECONDS=21600
# ===== 阿里云文档解析 =====
ALIBABA_ACCESS_KEY_ID=your_aliyun_access_key_id
@@ -96,11 +115,15 @@ RAG_TOP_K=10
RAG_RETRIEVAL_TOP_K=20
RAG_MAX_CONTEXT_TOKENS=4000
RAG_SUMMARY_MAX_TOKENS=1024
RAG_SKILLS_MAX_TOKENS=2048
# ===== Reranker配置(Cross-Encoder精排,默认关闭)=====
# 设置 RERANKER_ENABLED=true 并配置 RERANKER_BASE_URL 以启用精排
RERANKER_ENABLED=false
RERANKER_BASE_URL=
# ── Reranker (Cross-Encoder) ──────────────────────────────────────────────────
# Set RERANKER_ENABLED=true and point to a TEI or Cohere-compatible rerank API.
# Recommended model: BAAI/bge-reranker-v2.5-gemma2-lightweight (lighter) or
# BAAI/bge-reranker-v2-m3 (heavier, higher quality).
# The endpoint must expose POST /rerank (TEI style) or POST /v1/rerank (Cohere style).
RERANKER_ENABLED=true
RERANKER_BASE_URL=http://6.86.80.4:30080/v1
RERANKER_MODEL=BAAI/bge-reranker-v2-m3
RERANKER_API_KEY=
RERANKER_TOP_K=5
@@ -108,3 +131,61 @@ RERANKER_TOP_K=5
# ===== 会话配置 =====
SESSION_MAX_SESSIONS=100
SESSION_TIMEOUT_MINUTES=30
# SESSION_BACKEND=redis 启用 Redis 持久化会话(需要 Redis 可用,推荐生产环境)
# SESSION_BACKEND=memory 使用内存会话(重启丢失,适合本地开发)
SESSION_BACKEND=memory
# ===== 认证配置 (Auth) =====
# 生产环境必须替换为强随机密钥:
# python -c "import secrets; print(secrets.token_hex(32))"
AUTH_SECRET_KEY=change-me-in-production-must-be-32-or-more-characters-long
AUTH_ALGORITHM=HS256
# Token 有效期(分钟),默认 8 小时
AUTH_TOKEN_EXPIRE_MINUTES=480
# 设为 false 可跳过认证(仅限本地开发调试,生产必须 true)
AUTH_ENABLED=true
# ===== HyDE 查询增强 =====
# HyDE (Hypothetical Document Embeddings): 在检索前让 LLM 生成一段"假设性回答",
# 用该段落的 embedding 代替原始查询 embedding 进行向量检索。
# 无需新模型,复用现有 LLM 和 Embedding 服务。降低此功能可减少每次查询的 LLM 调用次数。
HYDE_ENABLED=true
HYDE_MAX_TOKENS=200
# ?????? LLM;???????????????
HYDE_LLM_PROVIDER=qwen
HYDE_LLM_MODEL=qwen3.6-flash
# ===== Agentic RAG 配置 (P0-1) =====
# 以下参数控制 /api/v1/agent/agentic/stream 多步推理管线
# 意图分类: simple_qa / compare / multi_hop / ambiguous
# compare 和 multi_hop 触发查询分解,最多 AGENTIC_MAX_SUB_QUERIES 个子查询
AGENTIC_MAX_SUB_QUERIES=4
# 引文锚定 fast-path 阈值: avg_score > 此值 且 chunks >= 3 时跳过 LLM grounding check
# 降低此值可让更多查询触发 LLM 二次验证(更准确,但延迟+成本增加)
AGENTIC_GROUNDING_THRESHOLD=0.65
# 各步骤 LLM 最大 token 数(越小越快,越大越准)
AGENTIC_INTENT_MAX_TOKENS=200
AGENTIC_PLAN_MAX_TOKENS=400
AGENTIC_GROUNDING_MAX_TOKENS=250
# ===== CORS =====
# 逗号分隔的允许跨域来源列表,生产环境绝不能使用 *
CORS_ALLOW_ORIGINS=http://localhost:5173
# ===== MCP (Model Context Protocol) =====
# MCP 端点(/mcp/)的 Host 头白名单,逗号分隔。MCP SDK 默认开启 DNS rebinding
# 防护,任何不在此列表中的 Host 都会被直接返回 HTTP 421,请求根本到不了鉴权和
# 工具逻辑。因此**远程部署必须把真实访问地址写进来**,否则所有外部 MCP 客户端
# Claude Desktop / IDE 等)100% 连不上。
# 语法:`:*` 后缀表示匹配任意端口;填 `*` 表示彻底关闭该防护(不推荐)。
# 例如部署在 6.86.80.9:8000 时:
# MCP_ALLOWED_HOSTS=6.86.80.9:*,127.0.0.1:*,localhost:*
MCP_ALLOWED_HOSTS=127.0.0.1:*,localhost:*,[::1]:*
# 系统状态页展示、以及"复制接入配置"按钮所使用的 MCP 外部访问地址。
# 留空则由后端从请求的 Host 头推导;当前端经 Vite 代理(changeOrigin: true
# 或反向代理改写了 Host 时,推导结果会是 127.0.0.1,此时必须显式指定。
# MCP_PUBLIC_URL=http://6.86.80.9:8000/mcp/
MCP_PUBLIC_URL=
+3
View File
@@ -62,3 +62,6 @@ logs/
# codex
.agents
# personal local records (never commit)
local/
@@ -0,0 +1,56 @@
<h2>Compliance Analysis — 哪个方向最值得优化?</h2>
<p class="subtitle">基于代码深度分析,发现了 4 个有价值的改进方向。选择你最希望深入的那个。</p>
<div class="options">
<div class="option" data-choice="A" onclick="toggleSelect(this)">
<div class="letter">A</div>
<div class="content">
<h3>⚡ 分析质量提升</h3>
<p>并行子句处理(速度 3–5×)、跨编码器重排序、置信度过滤、修复 highlight_terms 失效 Bug、减少 LLM 静默失败。</p>
<div class="pros-cons" style="margin-top:10px">
<div class="pros"><h4>收益</h4><ul><li>更快、更准确的分析</li><li>消除当前 Bug</li></ul></div>
<div class="cons"><h4>难度</h4><ul><li>需要改造 pipeline.py</li></ul></div>
</div>
</div>
</div>
<div class="option" data-choice="B" onclick="toggleSelect(this)">
<div class="letter">B</div>
<div class="content">
<h3>📋 分析历史 &amp; 专业报告</h3>
<p>持久化分析记录(PostgreSQL)、历史对比、PDF/DOCX 专业报告导出、分析版本追踪。</p>
<div class="pros-cons" style="margin-top:10px">
<div class="pros"><h4>收益</h4><ul><li>结果不再丢失</li><li>可交付给客户的报告</li></ul></div>
<div class="cons"><h4>难度</h4><ul><li>需要新增数据库表</li></ul></div>
</div>
</div>
</div>
<div class="option" data-choice="C" onclick="toggleSelect(this)">
<div class="letter">C</div>
<div class="content">
<h3>💬 深度 Chat 增强</h3>
<p>每个 Finding 独立对话线程(持久化)、Chat 上下文绑定真实检索到的法规原文、多轮追问记忆、快捷建议问句生成。</p>
<div class="pros-cons" style="margin-top:10px">
<div class="pros"><h4>收益</h4><ul><li>Finding 解读深度大幅提升</li><li>用户粘性强</li></ul></div>
<div class="cons"><h4>难度</h4><ul><li>需重构 chat 端点</li></ul></div>
</div>
</div>
</div>
<div class="option" data-choice="D" onclick="toggleSelect(this)">
<div class="letter">D</div>
<div class="content">
<h3>📑 自定义规则 &amp; 模板</h3>
<p>用户自定义合规规则库、按行业预设模板(汽车/金融/医疗)、Prompt 版本管理、A/B 测试不同提示策略。</p>
<div class="pros-cons" style="margin-top:10px">
<div class="pros"><h4>收益</h4><ul><li>适应不同行业场景</li><li>可配置,无需改代码</li></ul></div>
<div class="cons"><h4>难度</h4><ul><li>需要规则管理 UI</li></ul></div>
</div>
</div>
</div>
</div>
<p class="subtitle" style="margin-top:20px">💡 也可以多选,或者在终端告诉我你有其他想法。</p>
@@ -0,0 +1,3 @@
{"type":"click","text":"C\n \n 💬 深度 Chat 增强\n 每个 Finding 独立对话线程(持久化)、Chat 上下文绑定真实检索到的法规原文、多轮追问记忆、快捷建议问句生成。\n \n 收益Finding 解读深度大幅提升用户粘性强\n 难度需重构 chat 端点","choice":"C","id":null,"timestamp":1780897984866}
{"type":"click","text":"B\n \n 📋 分析历史 & 专业报告\n 持久化分析记录(PostgreSQL)、历史对比、PDF/DOCX 专业报告导出、分析版本追踪。\n \n 收益结果不再丢失可交付给客户的报告\n 难度需要新增数据库表","choice":"B","id":null,"timestamp":1780897985879}
{"type":"click","text":"A\n \n ⚡ 分析质量提升\n 并行子句处理(速度 3–5×)、跨编码器重排序、置信度过滤、修复 highlight_terms 失效 Bug、减少 LLM 静默失败。\n \n 收益更快、更准确的分析消除当前 Bug\n 难度需要改造 pipeline.py","choice":"A","id":null,"timestamp":1780897986554}
@@ -0,0 +1 @@
{"reason":"idle timeout","timestamp":1780894411095}
@@ -0,0 +1 @@
1055
+30 -4
View File
@@ -390,12 +390,38 @@ Demo-glm/
| 下载文档 | `/api/v1/documents/download/{doc_id}` | GET | 下载原文PDF/DOCX |
| 文档列表 | `/api/v1/documents/list` | GET | 列出已上传文档 |
| 检索知识 | `/api/v1/knowledge/search` | POST | 向量检索 |
| 单次问答 | `/api/v1/agent/ask` | POST | 智能问答 |
| 多轮对话 | `/api/v1/agent/chat` | POST | 会话对话 |
| 单次问答 | `/api/v1/agent/ask` | POST | 标准单轮问答 |
| 多轮对话 | `/api/v1/agent/chat` | POST | 标准会话对话 |
| 流式对话 | `/api/v1/agent/chat/stream` | POST | 标准流式问答 (SSE) |
| **Agentic 流式对话** | **`/api/v1/agent/agentic/stream`** | **POST** | **P0-1 多步推理 (SSE):意图分析→查询分解→迭代检索→引文锚定→生成** |
| 会话信息 | `/api/v1/agent/session/{id}` | GET | 获取会话 |
| 删除会话 | `/api/v1/agent/session/{id}` | DELETE | 删除会话 |
| Prompt模板 | `/api/v1/agent/templates` | GET | 模板列表 |
| 可用模型 | `/api/v1/agent/models` | GET | LLM模型列表 |
| 会话历史 | `/api/v1/agent/session/{id}/history` | GET | 获取历史记录 |
| 会话列表 | `/api/v1/agent/sessions` | GET | 列出所有会话 |
### Agentic 流式接口说明 (`/api/v1/agent/agentic/stream`)
**请求体** (同 `/agent/chat/stream`)
```json
{ "query": "GB 18384 与 ECE R100 在电池安全上有哪些差异?", "session_id": null, "top_k": 5 }
```
**额外 SSE 事件** (`thinking`)
```
event: thinking
data: {"step": "intent_analysis", "status": "done", "intent_type": "compare", "requires_decomposition": true}
event: thinking
data: {"step": "query_planning", "status": "done", "sub_queries": ["GB 18384 电池安全要求", "ECE R100 电池安全要求"]}
event: thinking
data: {"step": "retrieving", "status": "done", "query": "GB 18384 电池安全要求", "index": 1, "total": 2, "found": 8}
event: thinking
data: {"step": "grounding_check", "status": "done", "sufficient": true, "confidence": 0.82, "reason": "检索置信度充足"}
```
**意图类型**`simple_qa`(单跳)/ `compare`(对比)/ `multi_hop`(多跳)/ `ambiguous`(模糊)
---
+5
View File
@@ -0,0 +1,5 @@
"""FastAPI dependency functions for authentication and authorisation.
Import `get_current_user` or `require_role` into route modules to protect
endpoints. Both use the shared JWTHandler wired through bootstrap.
"""
+72
View File
@@ -0,0 +1,72 @@
"""FastAPI dependencies for JWT authentication.
Usage in a route:
from app.api.dependencies.auth import get_current_user, require_role
from app.domain.auth.models import UserRole
@router.get("/protected")
async def protected(user: UserClaims = Depends(get_current_user)):
return {"user": user.username}
@router.delete("/admin-only")
async def admin_only(user: UserClaims = Depends(require_role(UserRole.ADMIN))):
...
"""
from __future__ import annotations
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from app.config.settings import settings
from app.domain.auth.models import UserClaims, UserRole
from app.shared.bootstrap import get_jwt_handler
# Use Bearer token scheme — client sends `Authorization: Bearer <token>`.
_bearer = HTTPBearer(auto_error=False)
async def get_current_user(
credentials: HTTPAuthorizationCredentials | None = Depends(_bearer),
) -> UserClaims:
"""Extract and validate the JWT from the Authorization header.
Returns the decoded UserClaims on success.
Raises HTTP 401 when the token is missing, expired, or invalid.
When auth_enabled=False (development), returns a synthetic admin user.
"""
if not settings.auth_enabled:
# Development bypass — never enable this in production.
return UserClaims(user_id="dev", username="dev-admin", role=UserRole.ADMIN)
if credentials is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Missing authentication token",
headers={"WWW-Authenticate": "Bearer"},
)
try:
return get_jwt_handler().decode_token(credentials.credentials)
except ValueError as exc:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=str(exc),
headers={"WWW-Authenticate": "Bearer"},
) from exc
def require_role(*roles: UserRole):
"""Return a dependency that enforces one of the given roles.
Example:
Depends(require_role(UserRole.ADMIN, UserRole.LEGAL))
"""
async def _check(user: UserClaims = Depends(get_current_user)) -> UserClaims:
"""Verify the user holds one of the required roles."""
if user.role not in roles:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"Role '{user.role}' is not permitted. Required: {[r.value for r in roles]}",
)
return user
return _check
+27 -2
View File
@@ -1,6 +1,6 @@
"""FastAPI application entrypoint."""
from contextlib import asynccontextmanager
from contextlib import AsyncExitStack, asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.encoders import jsonable_encoder
@@ -8,10 +8,12 @@ from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from loguru import logger
from app.api.middleware.audit import AuditMiddleware
from app.api.models import ErrorResponse
from app.api.routes import api_router
from app.config.logging import setup_logging
from app.config.settings import settings
from app.mcp.server import build_mcp_asgi_app
from app.shared.bootstrap import cleanup_runtime_dependencies, preload_runtime_dependencies
from app.shared.errors import VectorStoreSchemaError
# Keep module behavior explicit so the backend flow stays easy to audit.
@@ -19,10 +21,23 @@ from app.shared.errors import VectorStoreSchemaError
setup_logging(level="INFO" if not settings.debug else "DEBUG")
# Built once at module scope so both lifespan() and app.mount() below reference
# the same instance — mounting a second, separately-built instance would start
# a second, unrelated MCP session manager.
mcp_app = build_mcp_asgi_app()
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Application lifecycle hooks."""
# FastMCP-style servers own a session manager that only starts via its own
# lifespan context. app.mount() does NOT propagate nested ASGI lifespans
# automatically (confirmed Starlette/ASGI limitation) — without this,
# every search_regulations call would fail because the MCP session
# manager was never started.
async with AsyncExitStack() as stack:
await stack.enter_async_context(mcp_app.router.lifespan_context(mcp_app))
logger.info(f"启动 {settings.app_name} v{settings.app_version}")
logger.info(f"调试模式: {settings.debug}")
logger.info("预加载LLM客户端...")
@@ -46,15 +61,25 @@ app = FastAPI(
redoc_url="/redoc",
)
# Tighten CORS — only allow configured origins.
# Set CORS_ALLOW_ORIGINS in .env to the real frontend URL in production.
_ORIGINS = [o.strip() for o in settings.cors_allow_origins.split(",") if o.strip()]
if not _ORIGINS:
_ORIGINS = ["http://localhost:5173"]
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_origins=_ORIGINS,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Audit middleware logs every authenticated API call for compliance traceability.
app.add_middleware(AuditMiddleware)
app.include_router(api_router, prefix="/api/v1")
app.mount("/mcp", mcp_app)
@app.exception_handler(VectorStoreSchemaError)
+1
View File
@@ -0,0 +1 @@
"""HTTP middleware for cross-cutting concerns: audit logging."""
+56
View File
@@ -0,0 +1,56 @@
"""Audit logging middleware.
Logs every API request with method, path, status code, response time,
and the authenticated user identity (extracted from the JWT when present).
Log lines are structured so they can be ingested by ELK / Loki.
"""
from __future__ import annotations
import time
from fastapi import Request, Response
from loguru import logger
from starlette.middleware.base import BaseHTTPMiddleware
class AuditMiddleware(BaseHTTPMiddleware):
"""Log all API calls. Skips health/docs paths to reduce noise."""
# Paths that produce no audit log entry.
_SKIP_PREFIXES = ("/health", "/docs", "/redoc", "/openapi.json")
async def dispatch(self, request: Request, call_next) -> Response:
"""Intercept the request, call the handler, and log the outcome."""
path = request.url.path
if path == "/" or any(path == p or path.startswith(p + "/") for p in self._SKIP_PREFIXES):
return await call_next(request)
start = time.perf_counter()
response = await call_next(request)
elapsed_ms = int((time.perf_counter() - start) * 1000)
# Extract user identity from JWT header for structured audit records.
# The token is not re-validated here — auth dependencies do that upstream.
user_id = "anonymous"
username = "anonymous"
auth_header = request.headers.get("authorization", "")
if auth_header.startswith("Bearer "):
try:
from app.shared.bootstrap import get_jwt_handler
claims = get_jwt_handler().decode_token(auth_header[7:])
user_id = claims.user_id
username = claims.username
except Exception:
pass
logger.info(
"AUDIT method={} path={} status={} elapsed_ms={} user_id={} username={}",
request.method,
path,
response.status_code,
elapsed_ms,
user_id,
username,
)
return response
+5
View File
@@ -42,6 +42,11 @@ class ChatRequest(BaseModel):
provider: Optional[str] = None
model: Optional[str] = None
top_k: Optional[int] = Field(default=None, ge=1, le=20)
# Optional document text uploaded by the user as conversation context.
# The text is injected directly into the LLM prompt so the model can
# answer questions about it without vector-store indexing.
context_text: Optional[str] = Field(default=None, max_length=12000)
context_filename: Optional[str] = Field(default=None, max_length=256)
class ChatResponse(BaseModel):
+4 -1
View File
@@ -1,6 +1,7 @@
"""Initialize the app.api.routes package."""
from fastapi import APIRouter
from .auth import router as auth_router
from .compliance import router as compliance_router
from .documents import router as documents_router
from .knowledge import router as knowledge_router
@@ -14,7 +15,8 @@ from .rag import router as rag_router
# Keep package boundaries explicit so backend imports stay predictable.
api_router = APIRouter()
# Keep package boundaries explicit so backend imports stay predictable.
# Auth routes first so /auth/token is easy to discover.
api_router.include_router(auth_router)
api_router.include_router(documents_router)
api_router.include_router(knowledge_router)
api_router.include_router(agent_router)
@@ -25,6 +27,7 @@ api_router.include_router(rag_router)
__all__ = [
"api_router",
"auth_router",
"documents_router",
"knowledge_router",
"agent_router",
+60 -1
View File
@@ -20,7 +20,11 @@ from app.api.models import (
)
from app.config.settings import settings
from app.shared.async_utils import iter_in_thread
from app.shared.bootstrap import get_agent_conversation_service, get_agent_session_service
from app.shared.bootstrap import (
get_agent_conversation_service,
get_agent_session_service,
get_agentic_conversation_service,
)
# Keep route handlers close to their transport-layer wiring for easier auditing.
@@ -182,3 +186,58 @@ async def submit_feedback(request: FeedbackRequest):
return {"message": "反馈已提交", "session_id": result.session_id, "message_index": result.message_index}
except ValueError as exc:
raise HTTPException(status_code=404, detail=str(exc))
# ── P0-1: Agentic RAG endpoint ────────────────────────────────────────────────
@router.post("/agentic/stream")
async def agentic_stream(request: ChatRequest):
"""Stream an Agentic RAG response with live multi-step reasoning trace.
Unlike the standard ``/chat/stream`` endpoint this route runs a full pipeline:
intent analysis → query planning → iterative retrieval → grounding check →
answer generation.
Extra SSE event types beyond the standard ones:
* ``thinking`` — reasoning sub-step progress; data is a JSON object with
``step`` (intent_analysis / query_planning / retrieving / grounding_check),
``status`` (running / done), and step-specific fields.
The ``sources``, ``content``, and ``done`` events are identical to the standard
chat-stream contract so the existing frontend parser can handle them without
changes.
"""
async def generate_sse() -> AsyncGenerator[str, None]:
"""Handle SSE generation for the agentic chat endpoint."""
try:
session_id_, event_stream = get_agentic_conversation_service().stream_agentic_chat(
query=request.query,
session_id=request.session_id,
filters=request.filters,
provider=request.provider or settings.llm_provider,
model=request.model or settings.llm_model,
top_k=request.top_k or settings.rag_top_k,
context_text=request.context_text,
context_filename=request.context_filename,
)
yield f"event: session\ndata: {json.dumps({'session_id': session_id_})}\n\n"
async for event_data in iter_in_thread(event_stream):
event_type = event_data.get("event", "content")
data = event_data.get("data", "")
if isinstance(data, (dict, list)):
yield f"event: {event_type}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
else:
yield f"event: {event_type}\ndata: {data}\n\n"
except Exception as exc:
yield f"event: error\ndata: {str(exc)}\n\n"
return StreamingResponse(
generate_sse(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
+63
View File
@@ -0,0 +1,63 @@
"""Authentication routes — token issuance only.
POST /auth/token — exchange username + password for a JWT.
GET /auth/me — return the current user identity (requires token).
"""
from __future__ import annotations
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordRequestForm
from pydantic import BaseModel
from app.api.dependencies.auth import get_current_user
from app.config.settings import settings
from app.domain.auth.models import UserClaims
from app.shared.bootstrap import get_jwt_handler, get_user_store
router = APIRouter(prefix="/auth", tags=["认证"])
class TokenResponse(BaseModel):
"""JWT token response body."""
access_token: str
token_type: str = "bearer"
expires_in: int
@router.post("/token", response_model=TokenResponse)
async def login(form: OAuth2PasswordRequestForm = Depends()):
"""Issue a JWT for valid username + password credentials.
Uses standard OAuth2 password grant form fields — compatible with
Swagger UI Authorize button.
"""
user = get_user_store().authenticate(form.username, form.password)
if user is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
headers={"WWW-Authenticate": "Bearer"},
)
token = get_jwt_handler().create_access_token(
user_id=user.id,
username=user.username,
role=user.role,
)
return TokenResponse(
access_token=token,
token_type="bearer",
expires_in=settings.auth_token_expire_minutes * 60,
)
@router.get("/me")
async def get_me(current_user: UserClaims = Depends(get_current_user)):
"""Return the identity of the currently authenticated user."""
return {
"user_id": current_user.user_id,
"username": current_user.username,
"role": current_user.role.value,
}
+297 -12
View File
@@ -7,10 +7,12 @@ import json
from pathlib import Path
from typing import AsyncGenerator, Optional
from fastapi import APIRouter, File, Form, UploadFile
from fastapi import APIRouter, Depends, File, Form, UploadFile
from fastapi.responses import StreamingResponse
from loguru import logger
from app.api.dependencies.auth import get_current_user
from app.domain.auth.models import UserClaims
from app.schemas.compliance import (
AnalyzeResponse,
ComplianceChatRequest,
@@ -75,6 +77,7 @@ async def analyze_stream(
file: Optional[UploadFile] = File(None),
domains: Optional[str] = Form(None),
title: Optional[str] = Form(None),
current_user: UserClaims = Depends(get_current_user),
):
"""Stream compliance analysis as SSE events.
@@ -82,10 +85,10 @@ async def analyze_stream(
Events: stage | source | finding | done | error
"""
from app.application.compliance.pipeline import (
check_clause_compliance,
detect_cross_clause_conflicts,
extract_text_from_doc_id,
extract_text_from_file,
retrieve_for_clause,
run_clauses_streaming,
split_into_clauses,
synthesize_conclusion,
)
@@ -133,22 +136,32 @@ async def analyze_stream(
await asyncio.sleep(0)
clauses: list[str] = await asyncio.to_thread(split_into_clauses, para_text, client)
# ── Stage 3: retrieve + gap check per clause ──────────────────
# ── Stage 3: progressive per-clause retrieve + gap check ──────
findings: list[dict] = []
total_clauses = len(clauses)
for i, clause in enumerate(clauses):
yield _sse({
"type": "stage",
"stage": "analyzing",
"label": f"Analyzing clause {i + 1}/{len(clauses)}",
"label": f"Analyzing {total_clauses} clauses",
})
# Emit initial progress so the frontend can show the total count
yield _sse({"type": "progress", "done": 0, "total": total_clauses})
await asyncio.sleep(0)
chunks = await asyncio.to_thread(
retrieve_for_clause, clause, retrieval_service, 5, domains or None
)
done_count = 0
# Stream results as each clause completes (not after all finish)
async for res in run_clauses_streaming(
clauses, retrieval_service, client,
top_k=5,
domains=domains or None,
):
done_count += 1
i = res["index"]
chunks = res["chunks"]
finding = res["finding"]
# Emit source events
# Emit source events for this clause
for chunk in chunks[:3]:
yield _sse({
"type": "source",
@@ -157,15 +170,25 @@ async def analyze_stream(
"score": round(float(getattr(chunk, "score", 0)), 3),
"status": "retrieved",
"full_content": (getattr(chunk, "text", "") or "")[:300],
"clause_index": i,
})
await asyncio.sleep(0)
finding = await asyncio.to_thread(check_clause_compliance, clause, chunks, client)
if finding:
findings.append(finding)
yield _sse({"type": "finding", **finding})
# Real progress update after each clause completes
yield _sse({"type": "progress", "done": done_count, "total": total_clauses})
await asyncio.sleep(0)
# ── Stage 3b: cross-clause conflict detection ─────────────────
if findings:
conflicts = await asyncio.to_thread(
detect_cross_clause_conflicts, findings, client
)
if conflicts:
yield _sse({"type": "conflicts", "items": conflicts})
# ── Stage 4: synthesize conclusion ────────────────────────────
yield _sse({"type": "stage", "stage": "concluding", "label": "Generating conclusion…"})
await asyncio.sleep(0)
@@ -175,6 +198,45 @@ async def analyze_stream(
)
yield _sse({"type": "done", **conclusion_data})
# Auto-save analysis to database
try:
from app.shared.bootstrap import get_compliance_repository
from app.domain.compliance.ports import AnalysisRecord, FindingRecord
from datetime import datetime
repo = get_compliance_repository()
finding_records = [
FindingRecord(
id="",
analysis_id="",
seq=i,
title=f.get("title", ""),
description=f.get("desc", ""),
status=f.get("status", "ok"),
clause_ref=f.get("clause_ref"),
)
for i, f in enumerate(findings)
]
record = AnalysisRecord(
id="",
created_at=datetime.utcnow(),
created_by=current_user.username if hasattr(current_user, "username") else None,
doc_name=file_name or (title or "Pasted text"),
standard_name=title or "",
risk_score=conclusion_data.get("risk_score", 0),
conclusion=conclusion_data.get("conclusion", ""),
actions=conclusion_data.get("actions", []),
para_text=conclusion_data.get("para_text", ""),
highlight_terms=conclusion_data.get("highlight_terms", []),
findings=finding_records,
)
analysis_id = await asyncio.to_thread(repo.save_analysis, record)
yield _sse({"type": "saved", "analysis_id": analysis_id})
except NotImplementedError:
pass # No postgres backend configured — skip saving
except Exception as exc:
logger.warning("Failed to auto-save compliance analysis: {}", exc)
except Exception as exc:
logger.exception("analyze-stream pipeline error")
yield _sse({"type": "error", "text": str(exc)})
@@ -222,3 +284,226 @@ async def compliance_chat(segment_id: int, request: ComplianceChatRequest):
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
)
@router.get("/history")
async def list_history(
limit: int = 20,
offset: int = 0,
current_user: UserClaims = Depends(get_current_user),
):
"""Return paginated list of saved compliance analyses (newest first)."""
from app.shared.bootstrap import get_compliance_repository
try:
repo = get_compliance_repository()
records = await asyncio.to_thread(repo.list_analyses, limit, offset)
return [
{
"id": r.id,
"created_at": r.created_at.isoformat(),
"created_by": r.created_by,
"doc_name": r.doc_name,
"standard_name": r.standard_name,
"risk_score": r.risk_score,
"finding_count": len(r.findings),
}
for r in records
]
except NotImplementedError:
return []
@router.get("/history/{analysis_id}")
async def get_history_item(
analysis_id: str,
current_user: UserClaims = Depends(get_current_user),
):
"""Return full analysis record including findings."""
from app.shared.bootstrap import get_compliance_repository
from fastapi import HTTPException
repo = get_compliance_repository()
record = await asyncio.to_thread(repo.get_analysis, analysis_id)
if not record:
raise HTTPException(status_code=404, detail="Analysis not found")
return {
"id": record.id,
"created_at": record.created_at.isoformat(),
"created_by": record.created_by,
"doc_name": record.doc_name,
"standard_name": record.standard_name,
"risk_score": record.risk_score,
"conclusion": record.conclusion,
"actions": record.actions,
"para_text": record.para_text,
"highlight_terms": record.highlight_terms,
"findings": [
{
"id": f.id,
"seq": f.seq,
"title": f.title,
"description": f.description,
"status": f.status,
"clause_ref": f.clause_ref,
}
for f in record.findings
],
}
@router.delete("/history/{analysis_id}", status_code=204)
async def delete_history_item(
analysis_id: str,
current_user: UserClaims = Depends(get_current_user),
):
"""Delete a saved analysis (cascade removes findings and chat messages)."""
from app.shared.bootstrap import get_compliance_repository
repo = get_compliance_repository()
await asyncio.to_thread(repo.delete_analysis, analysis_id)
@router.get("/history/{analysis_id}/download")
async def download_history_docx(
analysis_id: str,
current_user: UserClaims = Depends(get_current_user),
):
"""Return a DOCX compliance report for the given analysis."""
from app.shared.bootstrap import get_compliance_repository
from app.infrastructure.compliance.docx_export import generate_docx
from fastapi import HTTPException
from fastapi.responses import Response
repo = get_compliance_repository()
record = await asyncio.to_thread(repo.get_analysis, analysis_id)
if not record:
raise HTTPException(status_code=404, detail="Analysis not found")
docx_bytes = await asyncio.to_thread(generate_docx, record)
safe_name = (record.doc_name or "report").replace(" ", "_")[:50]
filename = f"compliance_{safe_name}_{record.created_at.strftime('%Y%m%d')}.docx"
return Response(
content=docx_bytes,
media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
)
@router.get("/analyses/{analysis_id}/findings/{finding_id}/chat")
async def get_finding_chat_history(
analysis_id: str,
finding_id: str,
current_user: UserClaims = Depends(get_current_user),
):
"""Return persisted chat messages for a finding thread, oldest first."""
from app.shared.bootstrap import get_compliance_repository
try:
repo = get_compliance_repository()
messages = await asyncio.to_thread(repo.get_messages, finding_id)
return messages
except NotImplementedError:
return []
@router.post("/analyses/{analysis_id}/findings/{finding_id}/suggestions")
async def get_finding_suggestions(
analysis_id: str,
finding_id: str,
current_user: UserClaims = Depends(get_current_user),
):
"""Generate 3 LLM-powered follow-up question suggestions for a finding."""
from app.application.compliance.pipeline import generate_suggestions
from app.shared.bootstrap import get_compliance_repository
from app.services.llm.llm_factory import get_llm_client
from fastapi import HTTPException
repo = get_compliance_repository()
analysis = await asyncio.to_thread(repo.get_analysis, analysis_id)
if not analysis:
raise HTTPException(status_code=404, detail="Analysis not found")
finding = next((f for f in analysis.findings if f.id == finding_id), None)
if not finding:
raise HTTPException(status_code=404, detail="Finding not found")
client = get_llm_client(provider=settings.llm_provider, model=settings.llm_model)
questions = await asyncio.to_thread(generate_suggestions, finding, analysis, client)
return {"questions": questions}
@router.post("/analyses/{analysis_id}/findings/{finding_id}/chat")
async def finding_chat(
analysis_id: str,
finding_id: str,
request: ComplianceChatRequest,
current_user: UserClaims = Depends(get_current_user),
):
"""Stream a grounded chat response for a specific finding.
Loads the finding and analysis from DB to build grounded context.
Persists both user message and assistant response to finding_chat_messages.
"""
from app.application.compliance.pipeline import build_finding_context
from app.shared.bootstrap import get_compliance_repository
from fastapi import HTTPException
repo = get_compliance_repository()
analysis = await asyncio.to_thread(repo.get_analysis, analysis_id)
if not analysis:
raise HTTPException(status_code=404, detail="Analysis not found")
finding = next((f for f in analysis.findings if f.id == finding_id), None)
if not finding:
raise HTTPException(status_code=404, detail="Finding not found")
# Persist user message
await asyncio.to_thread(
repo.save_message, analysis_id, finding_id, "user", request.query
)
# Build message history (last 10 messages = 5 turns)
history = await asyncio.to_thread(repo.get_messages, finding_id)
history_messages = [
{"role": m["role"], "content": m["content"]}
for m in history[-10:]
]
# Build grounded system context
system_context = build_finding_context(finding, analysis)
full_query = f"[Compliance Finding Context]\n{system_context}\n\nUser question: {request.query}"
assistant_buffer: list[str] = []
async def generate() -> AsyncGenerator[str, None]:
try:
_, event_stream = get_agent_conversation_service().stream_chat(
query=full_query,
top_k=5,
prompt_template="compliance_qa",
)
for event in event_stream:
event_type = event.get("event", "")
if event_type == "content":
text = event.get("data", "")
if text:
assistant_buffer.append(text)
yield _sse({"type": "chunk", "text": text})
elif event_type == "done":
yield _sse({"type": "done"})
await asyncio.sleep(0)
except Exception as exc:
logger.exception("finding_chat stream error")
yield _sse({"type": "error", "text": str(exc)})
finally:
# Persist assistant response after stream completes
full_response = "".join(assistant_buffer)
if full_response:
try:
await asyncio.to_thread(
repo.save_message, analysis_id, finding_id, "assistant", full_response
)
except Exception as exc:
logger.warning("Failed to persist assistant message: {}", exc)
return StreamingResponse(
generate(),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
)
+109 -5
View File
@@ -5,12 +5,15 @@ from __future__ import annotations
from io import BytesIO
from urllib.parse import quote
from fastapi import APIRouter, File, Form, HTTPException, UploadFile
from fastapi import APIRouter, BackgroundTasks, Depends, File, Form, HTTPException, UploadFile
from fastapi.responses import StreamingResponse
from loguru import logger
from app.api.dependencies.auth import get_current_user
from app.api.models import DocumentUploadResponse
from app.application.documents import DocumentProcessResult
from app.config.settings import settings
from app.domain.auth.models import UserClaims
from app.shared.bootstrap import get_document_command_service, get_document_query_service
# Keep route handlers close to their transport-layer wiring for easier auditing.
@@ -31,16 +34,60 @@ def _document_response(result: DocumentProcessResult) -> DocumentUploadResponse:
)
def _run_process_in_background(
*,
doc_id: str,
file_name: str,
final_doc_name: str,
content: bytes,
regulation_type: str,
version: str,
generate_summary: bool,
run_id: str | None,
) -> None:
"""Run document processing synchronously inside a FastAPI BackgroundTask thread.
FastAPI executes BackgroundTasks in a threadpool executor, so blocking I/O
(parser API calls, embedding, Milvus upsert) is safe here.
"""
try:
svc = get_document_command_service()
svc._process_document(
doc_id=doc_id,
file_name=file_name,
final_doc_name=final_doc_name,
content=content,
regulation_type=regulation_type,
version=version,
generate_summary=generate_summary,
run_id=run_id,
)
except Exception:
logger.exception("BackgroundTask document processing failed: doc_id={}", doc_id)
@router.post("/upload", response_model=DocumentUploadResponse)
async def upload_document(
background_tasks: BackgroundTasks,
file: UploadFile = File(..., description="上传的文档文件"),
doc_id: str | None = Form(None, description="客户端预分配的文档ID,不传则自动生成"),
doc_name: str | None = Form(None, description="文档名称"),
regulation_type: str | None = Form(None, description="法规类型"),
version: str | None = Form(None, description="文档版本"),
generate_summary: bool = Form(False, description="是否生成摘要"),
sync: bool = Form(False, description="同步处理(演示/测试用,默认异步处理)"),
current_user: UserClaims = Depends(get_current_user),
):
"""Handle upload document."""
"""Upload a document and process it asynchronously.
Default path (sync=false):
1. Store binary to MinIO immediately — returns within seconds.
2. Schedule parse→embed→index as a FastAPI BackgroundTask (same process,
threadpool) OR enqueue to Celery workers when USE_CELERY_WORKER=true.
3. Poll GET /documents/status/{doc_id} for progress.
sync=true path: full inline processing, blocks until complete (demo / CI use).
"""
content = await file.read()
if not file.filename:
raise HTTPException(status_code=400, detail="文件名不能为空")
@@ -48,7 +95,11 @@ async def upload_document(
raise HTTPException(status_code=400, detail="上传文件为空")
try:
result = get_document_command_service().upload_and_process(
svc = get_document_command_service()
if sync:
# Synchronous fallback: full inline processing.
result = svc.upload_and_process(
doc_id=doc_id,
file_name=file.filename,
content=content,
@@ -58,9 +109,59 @@ async def upload_document(
version=version or "",
generate_summary=generate_summary,
)
else:
# Step 1: store binary and create the document record (fast, sync).
stored_doc_id, run_id = svc.store_document(
doc_id=doc_id,
file_name=file.filename,
content=content,
content_type=file.content_type or "application/octet-stream",
doc_name=doc_name,
regulation_type=regulation_type or "",
version=version or "",
generate_summary=generate_summary,
)
final_doc_name = doc_name or file.filename
# Step 2: schedule processing via Celery worker OR FastAPI BackgroundTask.
if settings.use_celery_worker:
from app.infrastructure.tasks.document_tasks import process_document_task
process_document_task.delay(
doc_id=stored_doc_id,
file_name=file.filename,
doc_name=final_doc_name,
regulation_type=regulation_type or "",
version=version or "",
generate_summary=generate_summary,
run_id=run_id,
)
processing_note = "已入 Celery 队列,由 Worker 处理。"
else:
# Default: run in FastAPI's threadpool — no external worker needed.
background_tasks.add_task(
_run_process_in_background,
doc_id=stored_doc_id,
file_name=file.filename,
final_doc_name=final_doc_name,
content=content,
regulation_type=regulation_type or "",
version=version or "",
generate_summary=generate_summary,
run_id=run_id,
)
processing_note = "正在后台处理。"
result = DocumentProcessResult(
doc_id=stored_doc_id,
doc_name=final_doc_name,
status="stored",
message=f"文件已存储,{processing_note}请轮询 GET /documents/status/{{doc_id}} 查看进度。",
)
if result.status == "failed":
raise HTTPException(status_code=500, detail=result.message)
return _document_response(result)
except HTTPException:
raise
except Exception as exc:
@@ -106,7 +207,7 @@ async def download_document(doc_id: str):
@router.get("/list")
async def list_documents():
async def list_documents(current_user: UserClaims = Depends(get_current_user)):
"""List documents."""
documents = get_document_query_service().list_documents()
return {
@@ -140,6 +241,9 @@ async def get_document_management_list():
"updated_at": item.updated_at.isoformat(),
"regulation_type": item.regulation_type,
"version": item.version,
# True only when the original binary file is stored in MinIO.
# Milvus-only synthetic docs have no binary file — download is disabled.
"has_file": bool(item.object_name),
}
for item in documents
],
@@ -148,7 +252,7 @@ async def get_document_management_list():
@router.delete("/{doc_id}")
async def delete_document(doc_id: str):
async def delete_document(doc_id: str, current_user: UserClaims = Depends(get_current_user)):
"""Delete a document and its associated data."""
deleted = get_document_command_service().delete(doc_id)
if not deleted:
+106 -2
View File
@@ -4,10 +4,17 @@ from __future__ import annotations
import json
from fastapi import APIRouter, Query
from fastapi import APIRouter, Depends, Query
from fastapi.responses import StreamingResponse
from app.shared.bootstrap import get_perception_service
from app.shared.bootstrap import (
get_crawl_service,
get_event_store,
get_notification_store,
get_perception_service,
)
from app.api.dependencies.auth import get_current_user
from app.domain.auth.models import UserClaims
from app.shared.async_utils import iter_in_thread
router = APIRouter(prefix="/perception", tags=["智能感知"])
@@ -65,3 +72,100 @@ async def analyze_event(event_id: str):
"X-Accel-Buffering": "no",
},
)
@router.post("/crawl")
async def run_crawl(
body: dict = None,
current_user: UserClaims = Depends(get_current_user),
):
"""Trigger manual crawl of regulatory sources. Streams SSE progress.
Body (optional): {"sources": ["CATARC", "国标委·强制性", "EUR-Lex"]}
Omit sources to crawl all registered sources.
"""
sources: list[str] | None = (body or {}).get("sources")
crawl_svc = get_crawl_service()
async def crawl_stream():
async for item in iter_in_thread(crawl_svc.run_crawl(sources=sources)):
event_name = item.get("event", "message")
data = item.get("data", "")
if isinstance(data, (dict, list)):
data = json.dumps(data, ensure_ascii=False)
yield f"event: {event_name}\ndata: {data}\n\n"
return StreamingResponse(
crawl_stream(),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
)
@router.post("/events/{event_id}/process")
async def process_event(
event_id: str,
current_user: UserClaims = Depends(get_current_user),
):
"""Trigger LLM pipeline (extract + assess + diff) for a single event."""
from datetime import UTC, datetime
from app.infrastructure.perception.llm_pipeline import LlmPipeline
from app.shared.bootstrap import get_retrieval_service
event = get_perception_service().get_event(event_id)
if not event:
from fastapi import HTTPException
raise HTTPException(status_code=404, detail=f"Event {event_id} not found")
store = get_event_store()
pipeline = LlmPipeline()
structure = pipeline.extract_structure(event)
event.update(structure)
event["affected_docs"] = pipeline.assess_impact(event, get_retrieval_service())
event["processed_at"] = datetime.now(UTC).isoformat()
store.upsert(event)
return {"status": "ok", "event_id": event_id, "processed_at": event["processed_at"]}
@router.get("/events/{event_id}/diff")
async def get_event_diff(event_id: str):
"""Return semantic diff detail for an event (only available if previously crawled twice)."""
event = get_perception_service().get_event(event_id)
if not event:
from fastapi import HTTPException
raise HTTPException(status_code=404, detail=f"Event {event_id} not found")
if not event.get("change_summary"):
from fastapi import HTTPException
raise HTTPException(status_code=404, detail="No diff available for this event")
return {
"event_id": event_id,
"change_summary": event.get("change_summary"),
"changed_sections": event.get("changed_sections") or [],
"previous_hash": event.get("previous_hash"),
"content_hash": event.get("content_hash"),
}
@router.get("/notifications")
async def list_notifications(
limit: int = Query(default=20, ge=1, le=100),
current_user: UserClaims = Depends(get_current_user),
):
"""Return the newest in-app notifications plus this user's unread count.
Every logged-in user sees the same broadcast feed — there is no per-role
or per-topic subscription. "read" per item and the aggregate unread_count
both reflect only the calling user's own read receipts.
"""
store = get_notification_store()
items = store.list_for_user(current_user.user_id, limit=limit)
return {"items": items, "unread_count": store.unread_count(current_user.user_id)}
@router.post("/notifications/read")
async def mark_notifications_read(current_user: UserClaims = Depends(get_current_user)):
"""Mark every currently-unread notification read for the calling user."""
marked = get_notification_store().mark_all_read(current_user.user_id)
return {"marked": marked}
+88 -4
View File
@@ -3,16 +3,24 @@
from __future__ import annotations
import json
from typing import AsyncGenerator
import os
import re
import tempfile
from typing import AsyncGenerator, Optional
from fastapi import APIRouter
from fastapi import APIRouter, Depends, File, UploadFile
from fastapi.responses import StreamingResponse
from loguru import logger
from app.api.dependencies.auth import get_current_user
from app.config.settings import settings
from app.domain.auth.models import UserClaims
from app.schemas.rag import RagChatRequest, QuickQuestionsResponse, QuickQuestion
from app.shared.async_utils import iter_in_thread
from app.shared.bootstrap import get_agent_conversation_service
# Maximum characters of document text injected as LLM context (≈ 6 000 tokens).
_MAX_CONTEXT_CHARS = 8_000
router = APIRouter(prefix="/rag", tags=["RAG问答"])
@@ -26,14 +34,90 @@ _DEFAULT_QUICK_QUESTIONS = [
]
def _extract_text_from_bytes(content: bytes, filename: str) -> str:
"""Extract plain text from an uploaded file using the document parser.
Tries the configured parser first; falls back to raw UTF-8 decode for
plain-text formats (.txt, .md). Returns at most _MAX_CONTEXT_CHARS characters
so the text fits comfortably inside the LLM context window.
"""
suffix = os.path.splitext(filename or "doc.pdf")[1] or ".pdf"
# Fast path: plain-text files don't need a parser
if suffix.lower() in {".txt", ".md", ".csv"}:
try:
return content.decode("utf-8", errors="replace")[:_MAX_CONTEXT_CHARS]
except Exception:
pass
tmp_path = ""
try:
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
tmp.write(content)
tmp_path = tmp.name
from app.shared.bootstrap import get_document_command_service
svc = get_document_command_service()
parsed = svc.parser.parse(file_path=tmp_path, doc_id="ctx_extract", doc_name=filename)
if parsed.raw_text:
return parsed.raw_text[:_MAX_CONTEXT_CHARS]
# Fallback: join semantic blocks
return "\n".join(
b.get("text", "") for b in parsed.semantic_blocks if b.get("text")
)[:_MAX_CONTEXT_CHARS]
except Exception as exc:
logger.warning("Context text extraction failed for {}: {}", filename, exc)
return ""
finally:
if tmp_path:
try:
os.unlink(tmp_path)
except OSError:
pass
@router.post("/upload-context")
async def upload_context(
file: UploadFile = File(...),
current_user: UserClaims = Depends(get_current_user),
):
"""Extract text from an uploaded document and return it as conversation context.
The client stores the returned text and includes it in subsequent /rag/chat
requests via the context_text field — the LLM receives the document content
directly without requiring vector-store indexing.
"""
content = await file.read()
filename = file.filename or "document"
text = await __import__("asyncio").to_thread(_extract_text_from_bytes, content, filename)
if not text.strip():
from fastapi import HTTPException
raise HTTPException(status_code=422, detail="Could not extract text from the uploaded file.")
return {
"filename": filename,
"text": text,
"char_count": len(text),
"truncated": len(text) >= _MAX_CONTEXT_CHARS,
}
@router.post("/chat")
async def rag_chat(request: RagChatRequest):
"""Stream RAG Q&A using the real agent service."""
async def rag_chat(
request: RagChatRequest,
current_user: UserClaims = Depends(get_current_user),
):
"""Stream RAG Q&A using the real agent service.
When request.context_text is provided the document text is passed directly
to the answer generator as a dedicated document context section — RAG
retrieval still runs on the user's original question (not the document text)
so embedding quality is preserved for regulation chunk matching.
"""
session_id, event_stream = get_agent_conversation_service().stream_chat(
query=request.query,
session_id=request.session_id,
filters=request.filters,
top_k=request.top_k or settings.rag_top_k,
context_text=request.context_text,
context_filename=request.context_filename,
)
async def generate() -> AsyncGenerator[str, None]:
+185 -1
View File
@@ -1,18 +1,25 @@
"""Define API routes for status."""
import asyncio
import time
from typing import Any
from fastapi import APIRouter
from fastapi import APIRouter, Request
from app.config.settings import settings
from app.domain.retrieval import RetrievedChunk
from app.mcp.server import get_mcp_status
from app.services.llm.llm_factory import get_llm_client, get_llm_factory
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=["系统状态"])
@@ -23,6 +30,16 @@ _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():
@@ -111,3 +128,170 @@ async def get_health():
"max": settings.session_max_sessions,
},
}
def _normalize_llm_provider(raw_provider: str) -> str:
"""Normalize a raw LLM_PROVIDER/HYDE_LLM_PROVIDER settings string to the
canonical LLMProvider enum value, the SAME way LLMFactory.create() does.
TrackedLLMClient.chat() (tracked_client.py) always records usage under
`self._inner.config.provider.value` — the NORMALIZED enum value produced by
LLMFactory._parse_provider() — never the raw string a caller passed to
get_llm_client(). Reusing that same normalization here (instead of
duplicating the alias table) guarantees the tracker key this route reads
always agrees with the key TrackedLLMClient wrote, even when the raw
settings value is a non-canonical alias (e.g. "deepseek-v3") or different
casing. Falls back to the raw string, unchanged, if it does not match any
known provider/alias, so this passive status endpoint still renders
(as "never_called") instead of raising on a misconfigured provider string.
"""
try:
return get_llm_factory()._parse_provider(raw_provider).value
except ValueError:
return raw_provider
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 _normalize_llm_provider(settings.llm_provider), settings.llm_model
if role == "hyde_llm":
return (
_normalize_llm_provider(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"
elif role == "hyde_llm":
enabled = settings.hyde_enabled
if not enabled:
# Same "config always wins" override as the reranker branch above:
# report "disabled" even if HyDE ran successfully before being
# turned off in settings during 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`.
Skipped entirely for "hyde_llm" when settings.hyde_enabled is False,
mirroring _ping_reranker()'s disabled-skip pattern: when HyDE is turned
off (or reuses the main LLM, the default), issuing this ping would just be
a redundant duplicate chat call against the same model for no benefit.
"main_llm" is always pinged regardless of this check.
"""
if role == "hyde_llm" and not settings.hyde_enabled:
return
provider, model = _resolve_role_provider_model(role)
try:
client = get_llm_client(provider=provider, model=model)
except Exception as exc: # noqa: BLE001 - record, then re-raise so gather() still isolates this ping
# get_llm_client() can fail before any TrackedLLMClient exists to
# record the outcome itself (e.g. missing API key, unsupported
# provider string), so record the failure here directly, otherwise it
# would be invisible on the /status/models page afterward.
get_model_usage_tracker().record(provider=provider, model=model, success=False, error=str(exc))
raise
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]}
@router.get("/mcp")
async def get_mcp_server_status(request: Request):
"""Return MCP endpoint config, advertised tools, and per-tool call counters.
This route is a thin HTTP adapter: everything MCP-specific is assembled by
app.mcp.server.get_mcp_status(). The only thing decided here is the public
URL, because only the HTTP layer knows how the client reached us.
"""
# request.base_url already carries scheme/host/port and a trailing slash;
# strip it before appending so the result is ".../mcp/", not ".../mcp//".
public_url = settings.mcp_public_url or f"{str(request.base_url).rstrip('/')}/mcp/"
return await get_mcp_status(public_url)
+7 -1
View File
@@ -1,7 +1,13 @@
"""Initialize the app.application.agent package."""
from .services import AgentConversationService, AgentSessionFeedbackResult, AgentSessionService
from .agentic_service import AgenticConversationService
# Keep package boundaries explicit so backend imports stay predictable.
__all__ = ["AgentConversationService", "AgentSessionFeedbackResult", "AgentSessionService"]
__all__ = [
"AgentConversationService",
"AgentSessionFeedbackResult",
"AgentSessionService",
"AgenticConversationService",
]
@@ -0,0 +1,453 @@
"""Implement the Agentic RAG pipeline for multi-step reasoning (P0-1).
Architecture
------------
The pipeline adds four explicit reasoning steps before answer generation:
1. Intent Analysis — classify query type (simple_qa / compare / multi_hop / ambiguous)
2. Query Planning — for complex intents, decompose into focused sub-queries
3. Iterative Retrieval — retrieve for each sub-query, merge with deduplication
4. Grounding Check — verify retrieved context is sufficient; refine query when not
5. Answer Generation — stream final answer with citations (reuses AnswerGenerator)
Each step emits SSE ``thinking`` events so the frontend can render the live
reasoning trace. The pipeline is entirely synchronous and returns a generator so
it plugs into the same ``iter_in_thread`` pattern used by the existing chat routes.
"""
from __future__ import annotations
import json
from dataclasses import dataclass, field
from typing import Generator
from loguru import logger
from app.application.knowledge import KnowledgeRetrievalService
from app.application.agent.hyde_expander import HyDEExpander
from app.config.settings import settings
from app.domain.conversation import ConversationStore
from app.domain.retrieval import RetrievedChunk
from app.infrastructure.llm.openai_compatible_answer_generator import OpenAICompatibleAnswerGenerator
from app.services.llm.llm_factory import get_llm_client
# ── Prompts ───────────────────────────────────────────────────────────────────
# Each prompt is kept module-level for easy review and fine-tuning.
_INTENT_SYSTEM = (
"You are a query classifier for a Chinese regulatory compliance knowledge base.\n\n"
"Classify the query into exactly one of:\n"
'- "simple_qa" : Single-hop, factual question about one regulation or clause\n'
'- "compare" : Comparison between two or more regulations, standards, or versions\n'
'- "multi_hop" : Requires chaining facts across multiple regulations to answer\n'
'- "ambiguous" : Too vague or broad to retrieve effectively\n\n'
"Return ONLY valid JSON — no markdown, no extra text:\n"
'{"type": "...", "reason": "one sentence", "requires_decomposition": true/false}\n\n'
'"requires_decomposition" must be true for compare and multi_hop types.'
)
_PLAN_SYSTEM = (
"You are a query planner for a Chinese regulatory compliance knowledge base.\n\n"
"Decompose the query into 2-4 focused, self-contained sub-queries that together fully "
"address the original question. Each sub-query must target one specific regulation, "
"clause, or concept and be independently searchable.\n\n"
"Return ONLY a valid JSON array — no markdown, no extra text:\n"
'["sub-query 1", "sub-query 2", ...]'
)
_GROUNDING_SYSTEM = (
"You are a grounding verifier for a regulatory compliance QA system.\n\n"
"Given a query and retrieved regulation passages, decide whether the passages contain "
"sufficient, accurate information to answer the query.\n\n"
"Return ONLY valid JSON — no markdown, no extra text:\n"
'{"sufficient": true/false, "confidence": 0.0-1.0, "reason": "one sentence", '
'"refined_query": "a more specific search query if not sufficient, else null"}'
)
# ── Result dataclasses ────────────────────────────────────────────────────────
@dataclass
class IntentResult:
"""Capture the output of the intent-analysis step."""
type: str = "simple_qa"
reason: str = ""
requires_decomposition: bool = False
@dataclass
class GroundingResult:
"""Capture the output of the grounding-check step."""
sufficient: bool = True
confidence: float = 1.0
reason: str = ""
refined_query: str | None = None
# ── Service ───────────────────────────────────────────────────────────────────
class AgenticConversationService:
"""Multi-step Agentic RAG pipeline with live reasoning trace via SSE.
The service is intentionally synchronous so it can be wrapped in
``iter_in_thread`` by the route layer without any async boilerplate.
"""
def __init__(
self,
*,
retrieval_service: KnowledgeRetrievalService,
answer_generator: OpenAICompatibleAnswerGenerator,
conversation_store: ConversationStore,
) -> None:
"""Initialise with injected dependencies from the composition root."""
self.retrieval_service = retrieval_service
self.answer_generator = answer_generator
self.conversation_store = conversation_store
# HyDE expander is stateless — one instance shared for all requests.
self._hyde = HyDEExpander()
# ── Private helpers ───────────────────────────────────────────────────────
def _llm_json(
self,
system: str,
user: str,
provider: str | None,
model: str | None,
max_tokens: int = 300,
) -> dict | list | None:
"""Call the LLM with a JSON-only prompt and return the parsed result.
Returns ``None`` on any API or parse failure so callers can degrade
gracefully without raising.
"""
client = get_llm_client(
provider=provider or settings.llm_provider,
model=model or settings.llm_model,
)
resp = client.chat(
[{"role": "system", "content": system}, {"role": "user", "content": user}],
max_tokens=max_tokens,
temperature=0.1,
)
if not resp.is_success:
logger.warning("AgenticService LLM call failed: {}", resp.error)
return None
try:
raw = resp.content.strip()
# Strip accidental markdown code fences the model may add.
if raw.startswith("```"):
parts = raw.split("```")
raw = parts[1] if len(parts) > 1 else raw
if raw.startswith("json"):
raw = raw[4:]
return json.loads(raw.strip())
except (json.JSONDecodeError, IndexError) as exc:
logger.debug("AgenticService JSON parse failed: {} | raw={}", exc, resp.content[:200])
return None
def _analyze_intent(
self, query: str, provider: str | None, model: str | None
) -> IntentResult:
"""Classify query intent to select the appropriate retrieval strategy."""
data = self._llm_json(
_INTENT_SYSTEM,
f"Query: {query}",
provider,
model,
max_tokens=settings.agentic_intent_max_tokens,
)
if isinstance(data, dict):
return IntentResult(
type=str(data.get("type", "simple_qa")),
reason=str(data.get("reason", "")),
requires_decomposition=bool(data.get("requires_decomposition", False)),
)
return IntentResult(type="simple_qa", reason="fallback — classifier returned no JSON", requires_decomposition=False)
def _plan_queries(
self, query: str, intent_type: str, provider: str | None, model: str | None
) -> list[str]:
"""Decompose a complex query into focused, independently-retrievable sub-queries."""
data = self._llm_json(
_PLAN_SYSTEM,
f"Original query ({intent_type}): {query}",
provider,
model,
max_tokens=settings.agentic_plan_max_tokens,
)
if isinstance(data, list) and data:
# Cap at configured maximum to keep latency predictable.
return [str(q) for q in data[:settings.agentic_max_sub_queries] if q]
return [query]
def _check_grounding(
self,
query: str,
chunks: list[RetrievedChunk],
provider: str | None,
model: str | None,
) -> GroundingResult:
"""Verify whether retrieved chunks are sufficient to ground an accurate answer.
Uses a fast score-threshold heuristic first; falls back to an LLM call only
when scores are borderline so that the happy-path adds no extra latency.
"""
if not chunks:
return GroundingResult(
sufficient=False,
confidence=0.0,
reason="未检索到相关内容",
refined_query=None,
)
avg_score = sum(c.score for c in chunks) / len(chunks)
# Fast path: high-confidence retrieval → skip extra LLM call.
if avg_score > settings.agentic_grounding_threshold and len(chunks) >= 3:
return GroundingResult(
sufficient=True,
confidence=round(avg_score, 3),
reason="检索置信度充足,无需二次查询",
refined_query=None,
)
# LLM-based grounding check for borderline retrievals.
context_preview = "\n".join(
f"[{i + 1}] (score={c.score:.2f}) {c.text[:200]}" for i, c in enumerate(chunks[:5])
)
data = self._llm_json(
_GROUNDING_SYSTEM,
f"Query: {query}\n\nRetrieved passages:\n{context_preview}",
provider,
model,
max_tokens=settings.agentic_grounding_max_tokens,
)
if isinstance(data, dict):
return GroundingResult(
sufficient=bool(data.get("sufficient", True)),
confidence=float(data.get("confidence", 0.5)),
reason=str(data.get("reason", "")),
refined_query=data.get("refined_query") or None,
)
return GroundingResult(sufficient=True, confidence=0.5, reason="grounding check skipped (parse error)", refined_query=None)
@staticmethod
def _intent_to_template(intent_type: str) -> str:
"""Map an intent type to the best prompt template name for answer generation."""
mapping = {
"compare": "comparison",
"multi_hop": "compliance_qa",
"simple_qa": "compliance_qa",
"ambiguous": "compliance_qa",
}
return mapping.get(intent_type, "compliance_qa")
@staticmethod
def _deduplicate(chunks: list[RetrievedChunk], max_chunks: int) -> list[RetrievedChunk]:
"""Remove duplicate chunk IDs, preserving first-occurrence order up to max_chunks."""
seen: set[str] = set()
result: list[RetrievedChunk] = []
for chunk in chunks:
if chunk.chunk_id not in seen:
seen.add(chunk.chunk_id)
result.append(chunk)
if len(result) >= max_chunks:
break
return result
# ── Public interface ──────────────────────────────────────────────────────
def stream_agentic_chat(
self,
*,
query: str,
session_id: str | None = None,
filters: str | None = None,
provider: str | None = None,
model: str | None = None,
top_k: int = 5,
context_text: str | None = None,
context_filename: str | None = None,
) -> tuple[str, Generator[dict, None, None]]:
"""Run the full Agentic RAG pipeline and return ``(session_id, event_generator)``.
When context_text is provided (user-attached document) it is:
- Summarised and prepended to the intent-analysis prompt so the classifier
understands what kind of question is being asked.
- Treated as baseline grounding so the pipeline skips unnecessary retries
when the document itself is the primary source.
- Passed to the answer generator so the LLM sees the full document alongside
retrieved regulation chunks.
The generator yields SSE event dicts compatible with the route's
``iter_in_thread`` pattern.
"""
session = self.conversation_store.get_session(session_id) if session_id else None
if session is None:
session = self.conversation_store.create_session()
self.conversation_store.save_message(session.session_id, role="user", content=query)
history = [{"role": msg.role, "content": msg.content} for msg in session.messages[-10:]]
active_session_id = session.session_id
# Build a brief document summary for classifier/planner prompts (avoid
# passing the full text which could overwhelm small-context LLMs).
_doc_summary: str = ""
if context_text and context_text.strip():
_doc_label = context_filename or "document"
_preview = context_text.strip()[:400]
_doc_summary = f"[User has attached document: {_doc_label}]\nDocument preview: {_preview}\n\n"
def event_stream() -> Generator[dict, None, None]:
"""Execute all pipeline steps and yield SSE events."""
# ── Step 1: Intent Analysis ──────────────────────────────────────
yield {"event": "thinking", "data": {"step": "intent_analysis", "status": "running"}}
# Prepend doc summary so the classifier knows what the user is asking about
intent_user_msg = f"{_doc_summary}Query: {query}" if _doc_summary else f"Query: {query}"
data = self._llm_json(
_INTENT_SYSTEM, intent_user_msg, provider, model,
max_tokens=settings.agentic_intent_max_tokens,
)
if isinstance(data, dict):
intent = IntentResult(
type=str(data.get("type", "simple_qa")),
reason=str(data.get("reason", "")),
requires_decomposition=bool(data.get("requires_decomposition", False)),
)
else:
intent = IntentResult(type="simple_qa", reason="fallback", requires_decomposition=False)
logger.debug("Agentic intent: type={} decompose={}", intent.type, intent.requires_decomposition)
yield {
"event": "thinking",
"data": {
"step": "intent_analysis",
"status": "done",
"intent_type": intent.type,
"reason": intent.reason,
"requires_decomposition": intent.requires_decomposition,
},
}
# ── Step 2: Query Planning ───────────────────────────────────────
sub_queries: list[str] = [query]
if intent.requires_decomposition:
yield {"event": "thinking", "data": {"step": "query_planning", "status": "running"}}
plan_user_msg = f"{_doc_summary}Original query ({intent.type}): {query}" if _doc_summary else f"Original query ({intent.type}): {query}"
data_plan = self._llm_json(
_PLAN_SYSTEM, plan_user_msg, provider, model,
max_tokens=settings.agentic_plan_max_tokens,
)
if isinstance(data_plan, list) and data_plan:
sub_queries = [str(q) for q in data_plan[:settings.agentic_max_sub_queries] if q]
logger.debug("Agentic sub-queries ({}): {}", len(sub_queries), sub_queries)
yield {
"event": "thinking",
"data": {"step": "query_planning", "status": "done", "sub_queries": sub_queries},
}
# ── Step 3: Iterative Retrieval ──────────────────────────────────
# Always retrieve using the user's original question (NOT the document
# text) so embedding quality is preserved for regulation matching.
# HyDE enriches the retrieval query with a short hypothetical answer
# to close the vocabulary gap between terse queries and long documents.
candidate_k = max(top_k * 3, 15)
all_chunks: list[RetrievedChunk] = []
# For simple_qa with a single query, HyDE gives the biggest benefit
# (bridging vague/colloquial questions to formal document language).
# For compare/multi_hop, the planner already decomposed into precise
# sub-queries, so HyDE is less critical but still applied per sub-query.
for idx, sq in enumerate(sub_queries, start=1):
yield {
"event": "thinking",
"data": {"step": "retrieving", "status": "running", "query": sq, "index": idx, "total": len(sub_queries)},
}
# HyDE expansion: generate hypothetical answer, embed it for retrieval.
# Falls back to original sub-query if LLM call fails.
retrieval_query = self._hyde.expand(sq)
chunks = self.retrieval_service.retrieve(query=retrieval_query, top_k=candidate_k, filters=filters)
all_chunks.extend(chunks)
yield {
"event": "thinking",
"data": {"step": "retrieving", "status": "done", "query": sq, "index": idx, "total": len(sub_queries), "found": len(chunks)},
}
unique_chunks = self._deduplicate(all_chunks, max_chunks=top_k * 4)
# ── Step 4: Grounding Check ──────────────────────────────────────
yield {"event": "thinking", "data": {"step": "grounding_check", "status": "running"}}
# When the user has attached a document, the document itself provides
# baseline grounding — skip the re-query loop to avoid the LLM asking
# "please provide the document text" as a refined query.
if context_text and context_text.strip():
grounding = GroundingResult(
sufficient=True,
confidence=0.95,
reason="用户已附件上传文档,以文档内容为基础作答",
refined_query=None,
)
else:
grounding = self._check_grounding(query, unique_chunks, provider, model)
yield {
"event": "thinking",
"data": {
"step": "grounding_check",
"status": "done",
"sufficient": grounding.sufficient,
"confidence": grounding.confidence,
"reason": grounding.reason,
},
}
# Only retry from vector store when no document is attached and grounding failed
if not grounding.sufficient and grounding.refined_query and not context_text:
logger.info("Grounding insufficient — re-querying: {}", grounding.refined_query)
yield {
"event": "thinking",
"data": {"step": "retrieving", "status": "running", "query": grounding.refined_query, "index": 1, "total": 1, "retry": True},
}
# Apply HyDE to the refined query as well for better retrieval.
refined_hyde_query = self._hyde.expand(grounding.refined_query)
refined_chunks = self.retrieval_service.retrieve(query=refined_hyde_query, top_k=candidate_k, filters=filters)
all_chunks.extend(refined_chunks)
unique_chunks = self._deduplicate(all_chunks, max_chunks=top_k * 4)
yield {
"event": "thinking",
"data": {"step": "retrieving", "status": "done", "query": grounding.refined_query, "index": 1, "total": 1, "found": len(refined_chunks), "retry": True},
}
final_chunks = unique_chunks[:top_k]
# ── Step 5: Answer Generation ────────────────────────────────────
sources_payload = [s.__dict__ for s in self.answer_generator._sources(final_chunks)]
yield {"event": "sources", "data": sources_payload}
answer_parts: list[str] = []
for event in self.answer_generator.stream_generate(
query=query,
retrieved_chunks=final_chunks,
history=history,
provider=provider,
model=model,
prompt_template=self._intent_to_template(intent.type),
context_text=context_text,
context_filename=context_filename,
):
if event.get("event") == "content":
answer_parts.append(str(event.get("data", "")))
yield event
full_answer = "".join(answer_parts)
self.conversation_store.save_message(
active_session_id,
role="assistant",
content=full_answer,
sources=sources_payload,
)
return active_session_id, event_stream()
@@ -0,0 +1,105 @@
"""Implement HyDE (Hypothetical Document Embeddings) query expansion.
HyDE improves dense retrieval by addressing the vocabulary gap between
short user queries and longer document passages:
User query → [LLM generates hypothetical answer]
embed hypothetical answer (not original query)
retrieve similar real passages from Milvus
The hypothetical answer uses the same vocabulary and phrasing as documents,
so its embedding is much closer to relevant chunks than a terse query embedding.
Usage:
expander = HyDEExpander()
retrieval_query = expander.expand(query, provider=..., model=...)
chunks = retrieval_service.retrieve(query=retrieval_query, ...)
When the LLM call fails, expand() falls back to the original query so the
retrieval pipeline degrades gracefully.
References:
Gao et al. (2022), "Precise Zero-Shot Dense Retrieval without Relevance Labels"
https://arxiv.org/abs/2212.10496
"""
from __future__ import annotations
from loguru import logger
from app.config.settings import settings
from app.services.llm.llm_factory import get_llm_client
# Maximum chars to trim from the hypothetical answer to avoid token overrun.
_MAX_HYPOTHESIS_CHARS = 600
# System prompt that instructs the LLM to write a passage *as if* it were
# from a regulatory document, not a conversation answer.
_HYDE_SYSTEM = (
"你是一位法规知识库专家。用户提出了一个问题,"
"请用50-120字写一段话,模拟如果相关法规文档中存在完美答案,"
"该段落会是什么内容。\n\n"
"要求:\n"
"- 使用与法规文档相同的正式书面语气\n"
"- 包含可能的条款编号、标准名称等关键术语\n"
"- 不要解释你在做什么,直接输出假设性段落\n"
"- 如问题过于模糊,写一段合理的通用法规说明"
)
class HyDEExpander:
"""Generate a hypothetical document passage to improve dense retrieval.
The expander is stateless — instantiate once and call expand() per query.
It requires no external dependencies beyond the project's existing LLM
client infrastructure.
"""
def expand(self, query: str) -> str:
"""Return a combined retrieval query: original query + hypothetical passage.
The combination ensures:
- Dense retrieval uses the enriched hypothetical text (semantic match).
- BM25 retrieval still benefits from the original query keywords.
The model used is ``settings.hyde_llm_model`` (dedicated lightweight model)
falling back to the main ``settings.llm_model`` when not configured.
If the LLM call fails for any reason, returns the original query unchanged.
"""
if not settings.hyde_enabled:
return query
# Use the dedicated HyDE model when configured; fall back to main LLM.
# A lightweight model (e.g. qwen3.6-flash) is sufficient for generating
# a short hypothetical passage and significantly reduces cost + latency.
provider = settings.hyde_llm_provider or settings.llm_provider
model = settings.hyde_llm_model or settings.llm_model
try:
client = get_llm_client(provider=provider, model=model)
resp = client.chat(
messages=[
{"role": "system", "content": _HYDE_SYSTEM},
{"role": "user", "content": f"问题:{query}"},
],
max_tokens=settings.hyde_max_tokens,
# Low temperature: we want a plausible, deterministic passage.
temperature=0.3,
)
if not resp.is_success or not resp.content:
logger.debug("HyDE LLM call failed or empty — using original query")
return query
hypothesis = resp.content.strip()[:_MAX_HYPOTHESIS_CHARS]
logger.debug("HyDE expanded query ({}{} chars)", len(query), len(hypothesis))
# Concatenate: the embedding model will see the full combined text,
# so the resulting vector leans toward the hypothetical document style.
return f"{query}\n\n{hypothesis}"
except Exception as exc: # noqa: BLE001 — intentional broad catch for graceful fallback
logger.warning("HyDE expansion failed: {} — using original query", exc)
return query
+20 -2
View File
@@ -9,6 +9,7 @@ from app.domain.conversation import AnswerGenerator, AnswerResult, ConversationS
from app.domain.retrieval import RetrievedChunk
from app.application.knowledge import KnowledgeRetrievalService
from app.application.agent.hyde_expander import HyDEExpander
# Keep orchestration logic centralized so use-case flow stays easy to trace.
@@ -26,6 +27,8 @@ class AgentConversationService:
self.retrieval_service = retrieval_service
self.answer_generator = answer_generator
self.conversation_store = conversation_store
# Shared HyDE expander — stateless, safe for reuse across requests.
self._hyde = HyDEExpander()
def ask(
self,
@@ -108,14 +111,26 @@ class AgentConversationService:
model: str | None = None,
top_k: int = 5,
prompt_template: str | None = None,
context_text: str | None = None,
context_filename: str | None = None,
) -> tuple[str, Generator[dict, None, None]]:
"""Stream chat for the Agent Conversation Service instance."""
"""Stream chat for the Agent Conversation Service instance.
When context_text is provided the user's document is passed directly to
the answer generator — RAG retrieval still runs on the user's question
(not the document text) to find relevant regulation passages.
"""
session = self.conversation_store.get_session(session_id) if session_id else None
if session is None:
session = self.conversation_store.create_session()
self.conversation_store.save_message(session.session_id, role="user", content=query)
history = [{"role": msg.role, "content": msg.content} for msg in session.messages[-10:]]
retrieved = self.retrieval_service.retrieve(query=query, top_k=top_k, filters=filters)
# HyDE: expand the query with a hypothetical answer to improve dense retrieval.
# For document-context queries, skip HyDE since the document itself guides retrieval.
retrieval_query = self._hyde.expand(query) if not context_text else query
# Retrieve using the enriched query — NOT the document text —
# so embedding quality is preserved for regulation chunk matching.
retrieved = self.retrieval_service.retrieve(query=retrieval_query, top_k=top_k, filters=filters)
def event_stream() -> Generator[dict, None, None]:
"""Handle event stream for the Agent Conversation Service instance."""
@@ -129,6 +144,8 @@ class AgentConversationService:
provider=provider,
model=model,
prompt_template=prompt_template,
context_text=context_text,
context_filename=context_filename,
):
if event.get("event") == "sources":
sources_payload = event.get("data", [])
@@ -189,3 +206,4 @@ class AgentSessionService:
raise ValueError("消息索引不存在")
# Preserve the existing API behavior until a persistent feedback store is introduced.
return AgentSessionFeedbackResult(session_id=session_id, message_index=message_index)
+384 -29
View File
@@ -5,6 +5,7 @@ All functions are synchronous — call them via asyncio.to_thread() in async SSE
from __future__ import annotations
import asyncio
import json
import os
import re
@@ -12,10 +13,20 @@ import tempfile
from typing import TYPE_CHECKING
from loguru import logger
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
# Shared retry policy for LLM calls: 3 attempts, exponential back-off 14 s.
_llm_retry = retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=1, max=4),
retry=retry_if_exception_type((ValueError, TimeoutError, ConnectionError)),
reraise=True,
)
if TYPE_CHECKING:
from app.application.knowledge import KnowledgeRetrievalService
from app.domain.retrieval import RetrievedChunk
from app.domain.compliance.ports import AnalysisRecord, FindingRecord
from app.services.llm.base_client import BaseLLMClient
@@ -40,19 +51,36 @@ def _extract_json(text: str):
def extract_text_from_doc_id(doc_id: str) -> str:
"""Fetch the full text of a document by retrieving its chunks filtered by doc_id.
Uses a high top_k and doc_id filter to reconstruct the document in chunk order,
avoiding the previous approach of semantic search by doc_name which could return
chunks from unrelated documents.
"""
from app.shared.bootstrap import get_document_query_service, get_retrieval_service
doc = get_document_query_service().get(doc_id)
if not doc:
raise ValueError(f"Document '{doc_id}' not found")
service = get_retrieval_service()
chunks = service.retrieve(query=doc.doc_name, top_k=30)
doc_chunks = [c for c in chunks if c.doc_id == doc_id]
# Use doc_name as a broad query, filter strictly by doc_id so we only get
# this document's chunks; top_k=100 covers most real-world documents.
chunks = service.retrieve(query=doc.doc_name, top_k=100, filters=doc_id)
doc_chunks = [c for c in chunks if getattr(c, "doc_id", None) == doc_id]
if not doc_chunks:
doc_chunks = chunks[:15]
return "\n\n".join(c.text for c in doc_chunks[:15])
# Fallback: use top results even without doc_id match (e.g., legacy store)
doc_chunks = chunks[:30]
# Sort by chunk_index to preserve document reading order
doc_chunks.sort(key=lambda c: getattr(c, "chunk_index", 0))
return "\n\n".join(c.text for c in doc_chunks[:40])
def extract_text_from_file(content: bytes, filename: str) -> str:
"""Parse an uploaded file and return its full text content.
Removed previous 4000-char cap so large specifications and standards are
fully analysed. The caller is responsible for splitting the text into
clause-sized chunks before passing to the LLM.
"""
from app.shared.bootstrap import get_document_command_service
suffix = os.path.splitext(filename or "doc.pdf")[1] or ".pdf"
tmp_path = ""
@@ -63,10 +91,11 @@ def extract_text_from_file(content: bytes, filename: str) -> str:
service = get_document_command_service()
parsed = service.parser.parse(file_path=tmp_path, doc_id="tmp_analysis", doc_name=filename)
if parsed.raw_text:
return parsed.raw_text[:4000]
# Return full text — truncation happens in split_into_clauses()
return parsed.raw_text
return "\n".join(
b.get("text", "") for b in parsed.semantic_blocks[:30] if b.get("text")
)[:4000]
b.get("text", "") for b in parsed.semantic_blocks if b.get("text")
)
except Exception as exc:
logger.warning("File text extraction failed: {}", exc)
return ""
@@ -77,27 +106,68 @@ def extract_text_from_file(content: bytes, filename: str) -> str:
def split_into_clauses(text: str, client: "BaseLLMClient") -> list[str]:
"""Split a compliance document into semantically independent clauses.
For long texts (> 2 000 chars) the document is processed in overlapping
2 000-char windows so no content is missed. Each window produces up to 4
clauses; results are deduplicated and capped at 12 total to keep analysis
latency reasonable.
"""
# Window size and step for sliding-window clause extraction
_WINDOW = 2000
_STEP = 1800 # 200-char overlap to avoid cutting clauses at boundaries
_MAX_CLAUSES = 12
windows = []
if len(text) <= _WINDOW:
windows = [text]
else:
pos = 0
while pos < len(text):
windows.append(text[pos: pos + _WINDOW])
pos += _STEP
all_clauses: list[str] = []
for window in windows:
prompt = (
"You are a compliance analysis expert. Split the following text into 3-8 "
"semantically complete compliance clauses. Each clause should be an independent "
"compliance requirement or technical statement.\n"
"You are a compliance analysis expert. Split the following text into "
"3-4 semantically complete compliance clauses. Each clause must be an "
"independent requirement or technical statement. Omit section headings, "
"definitions, and non-normative text.\n"
"Return as JSON array of strings, e.g.:\n"
'["Clause one...", "Clause two..."]\n'
"Return ONLY the JSON array.\n\n"
f"Text:\n{text[:2000]}"
f"Text:\n{window}"
)
response = client.chat([{"role": "user", "content": prompt}], max_tokens=1000)
response = client.chat([{"role": "user", "content": prompt}], max_tokens=800)
if response.is_success:
try:
result = _extract_json(response.content)
if isinstance(result, list):
clauses = [str(c).strip() for c in result if str(c).strip()]
if clauses:
return clauses[:8]
all_clauses.extend(clauses[:4])
except (ValueError, TypeError):
logger.warning("Clause split JSON parse failed, using fallback")
sentences = re.split(r"[.?!;\n]+", text)
return [s.strip() for s in sentences if len(s.strip()) > 20][:6]
logger.warning("Clause split JSON parse failed for window, using sentence fallback")
sentences = re.split(r"[.?!;\n]+", window)
all_clauses.extend(s.strip() for s in sentences if len(s.strip()) > 20)
else:
# LLM unavailable — fall back to sentence splitting for this window
sentences = re.split(r"[.?!;\n]+", window)
all_clauses.extend(s.strip() for s in sentences if len(s.strip()) > 20)
if len(all_clauses) >= _MAX_CLAUSES:
break
# Deduplicate near-duplicates (same first 80 chars) that span window boundaries
seen: set[str] = set()
deduped: list[str] = []
for c in all_clauses:
key = c[:80].lower()
if key not in seen:
seen.add(key)
deduped.append(c)
return deduped[:_MAX_CLAUSES]
def retrieve_for_clause(
@@ -106,7 +176,130 @@ def retrieve_for_clause(
top_k: int = 5,
domains: str | None = None,
) -> list["RetrievedChunk"]:
return retrieval_service.retrieve(query=clause, top_k=top_k, filters=domains)
"""Retrieve regulation chunks relevant to a clause.
If the best retrieval score is below 0.55, rewrite the clause into a more
technical query and retry once to improve coverage.
"""
chunks = retrieval_service.retrieve(query=clause, top_k=top_k, filters=domains)
if not chunks:
return chunks
best_score = max((getattr(c, "score", 0) for c in chunks), default=0)
if best_score < 0.55:
# Rewrite clause as technical keyword query and retry
keywords = " ".join(
w for w in re.split(r"\W+", clause) if len(w) > 3
)[:200]
retry_chunks = retrieval_service.retrieve(query=keywords, top_k=top_k, filters=domains)
if retry_chunks:
# Merge: keep unique chunks, prefer higher-score version
seen_ids: set[str] = {getattr(c, "chunk_id", str(i)) for i, c in enumerate(chunks)}
for rc in retry_chunks:
rid = getattr(rc, "chunk_id", "")
if rid not in seen_ids:
chunks.append(rc)
seen_ids.add(rid)
chunks.sort(key=lambda c: getattr(c, "score", 0), reverse=True)
chunks = chunks[:top_k]
return chunks
def process_single_clause(
clause: str,
index: int,
retrieval_service: "KnowledgeRetrievalService",
client: "BaseLLMClient",
top_k: int = 5,
domains: str | None = None,
) -> dict:
"""Process one clause: retrieve relevant regulations then check compliance.
Returns a dict with keys:
- index: clause position (for ordering)
- chunks: list of RetrievedChunk (for source events)
- finding: dict with title/desc/status/clause_ref/confidence (may be None on LLM failure)
Designed to run inside asyncio.to_thread() for parallel execution.
The finding now includes a 'source_refs' list linking back to the chunks
that informed the verdict, enabling the frontend to correlate sources with findings.
"""
chunks = retrieve_for_clause(clause, retrieval_service, top_k, domains)
finding = check_clause_compliance(clause, chunks, client)
if finding is not None:
# Attach source references so the frontend can link finding ↔ sources
finding["source_refs"] = [
{
"standard": getattr(c, "doc_title", "") or getattr(c, "doc_name", ""),
"clause": getattr(c, "section_title", "") or "",
"score": round(float(getattr(c, "score", 0)), 3),
}
for c in chunks[:3]
]
return {"index": index, "chunks": chunks, "finding": finding}
async def run_clauses_streaming(
clauses: list[str],
retrieval_service: "KnowledgeRetrievalService",
client: "BaseLLMClient",
top_k: int = 5,
domains: str | None = None,
):
"""Process all clauses concurrently and yield each result as it completes.
Unlike the old gather()-based approach, this uses asyncio.Queue so that
findings are emitted to the SSE stream immediately when each clause
finishes — the user sees results progressively rather than waiting for
the slowest clause before seeing any output.
Yields dicts with keys: index, chunks, finding (same schema as
process_single_clause, plus a sentinel {"_done": True} at the end).
"""
queue: asyncio.Queue[dict] = asyncio.Queue()
total = len(clauses)
async def _worker(clause: str, i: int) -> None:
"""Run one clause in a thread and push the result into the queue."""
try:
result = await asyncio.to_thread(
process_single_clause,
clause, i, retrieval_service, client, top_k, domains,
)
except Exception as exc:
logger.warning("Clause {} processing failed: {}", i, exc)
result = {"index": i, "chunks": [], "finding": None}
await queue.put(result)
# Launch all workers concurrently
tasks = [asyncio.create_task(_worker(clause, i)) for i, clause in enumerate(clauses)]
received = 0
while received < total:
result = await queue.get()
yield result
received += 1
# Wait for all tasks to complete (they should already be done by now)
await asyncio.gather(*tasks, return_exceptions=True)
async def run_clauses_parallel(
clauses: list[str],
retrieval_service: "KnowledgeRetrievalService",
client: "BaseLLMClient",
top_k: int = 5,
domains: str | None = None,
) -> list[dict]:
"""Legacy batch API kept for backward compatibility.
Collects all streaming results and returns them sorted by clause index.
New code should use run_clauses_streaming() directly.
"""
results: list[dict] = []
async for result in run_clauses_streaming(clauses, retrieval_service, client, top_k, domains):
results.append(result)
return sorted(results, key=lambda r: r["index"])
def check_clause_compliance(
@@ -114,30 +307,50 @@ def check_clause_compliance(
chunks: list["RetrievedChunk"],
client: "BaseLLMClient",
) -> dict | None:
if not chunks:
return None
"""Check whether a business clause complies with the retrieved regulations.
The prompt explicitly instructs the LLM to:
- extract clause_ref from the retrieved text (not invent it)
- include a confidence score (0-1) reflecting how well the retrieved
chunks cover the clause topic
Returns None only when the LLM call fails after all retries.
"""
reg_context = "\n".join(
f"[{i+1}] {c.doc_title} {c.section_title or ''}: {c.text[:300]}"
for i, c in enumerate(chunks[:5])
)
) if chunks else "(no regulatory context retrieved)"
prompt = (
"You are a compliance expert. Judge whether the following business clause "
"complies with the retrieved regulations.\n\n"
f"Business clause:\n{clause}\n\n"
f"Retrieved regulations:\n{reg_context}\n\n"
"Return JSON:\n"
"Return JSON with these exact fields:\n"
"{\n"
' "status": "ok" | "warn" | "risk",\n'
' "title": "Short finding title (max 30 chars)",\n'
' "desc": "Description (50-120 chars)",\n'
' "clause_ref": "Regulation clause reference e.g. Art.9.1 or Sec.3.1"\n'
' "clause_ref": "Exact clause/article reference copied from the retrieved text above, '
'e.g. Art.9.1 or Sec.3.1. Use null if no specific clause number appears in the retrieved text.",\n'
' "confidence": 0.0-1.0 // how well the retrieved context covers this clause topic\n'
"}\n"
"status: ok=compliant, warn=gap exists, risk=critical/missing\n"
"IMPORTANT: copy clause_ref verbatim from the retrieved text; do NOT invent references.\n"
"Return ONLY the JSON object."
)
response = client.chat([{"role": "user", "content": prompt}], max_tokens=500)
if not response.is_success:
def _do_check():
resp = client.chat([{"role": "user", "content": prompt}], max_tokens=500)
if not resp.is_success:
raise ValueError("LLM returned non-success for gap check")
return resp
try:
response = _llm_retry(_do_check)()
except Exception as exc:
logger.warning("check_clause_compliance LLM call failed after retries: {}", exc)
return None
try:
result = _extract_json(response.content)
if isinstance(result, dict) and "status" in result:
@@ -145,7 +358,10 @@ def check_clause_compliance(
"title": str(result.get("title", "Compliance finding")),
"desc": str(result.get("desc", "")),
"status": result.get("status", "info"),
"clause_ref": result.get("clause_ref"),
# None if LLM correctly found no clause number in retrieved text
"clause_ref": result.get("clause_ref") or None,
# Confidence score helps frontend show retrieval quality indicator
"confidence": float(result.get("confidence", 0.5)),
}
except (ValueError, TypeError) as exc:
logger.warning("Gap check JSON parse failed: {}", exc)
@@ -182,12 +398,11 @@ def synthesize_conclusion(
' {"label": "Priority", "value": "High/Medium/Low", "risk": true}\n'
' ],\n'
' "risk_score": 0-100 (integer, higher=riskier),\n'
' "highlight_terms": ["Key terms to highlight, max 10 terms"],\n'
' "highlight_terms": ["term1", "term2"], // up to 10 key technical/legal terms actually present in the text\n'
' "para_text": "Original text or summary (max 600 chars)"\n'
"}\n"
"Return ONLY the JSON object."
)
response = client.chat([{"role": "user", "content": prompt}], max_tokens=1200)
fallback = {
"conclusion": "Compliance analysis complete. Review findings and create remediation plan.",
"actions": [
@@ -198,8 +413,19 @@ def synthesize_conclusion(
"highlight_terms": [],
"para_text": para_text[:800],
}
if not response.is_success:
def _do_synthesize():
resp = client.chat([{"role": "user", "content": prompt}], max_tokens=1200)
if not resp.is_success:
raise ValueError("LLM returned non-success for synthesis")
return resp
try:
response = _llm_retry(_do_synthesize)()
except Exception as exc:
logger.warning("synthesize_conclusion LLM call failed after retries: {}", exc)
return fallback
try:
result = _extract_json(response.content)
if isinstance(result, dict):
@@ -213,3 +439,132 @@ def synthesize_conclusion(
except (ValueError, TypeError) as exc:
logger.warning("Conclusion synthesis JSON parse failed: {}", exc)
return fallback
_SUGGESTION_FOCUS = {
"risk": "Focus on remediation steps, required certifications, and timeline to resolve.",
"warn": "Focus on identifying the specific compliance gap and how to close it.",
"ok": "Focus on maintaining compliance evidence and monitoring future changes.",
}
_SUGGESTION_FALLBACK = {
"risk": [
"What specific certifications or documents are required to remediate this finding?",
"What is the typical remediation timeline for this type of non-compliance?",
"Which regulation clause defines the exact requirement?",
],
"warn": [
"What is the exact gap between the current state and the requirement?",
"What evidence would demonstrate partial compliance?",
"Which regulation clause applies to this warning?",
],
"ok": [
"What documentation should be maintained to evidence this compliance?",
"How should this area be monitored as regulations evolve?",
"Are there related clauses that may affect this compliant area?",
],
}
def build_finding_context(finding: "FindingRecord", analysis: "AnalysisRecord") -> str:
"""Build a grounded system context string for a finding chat thread.
Combines finding details with analysis metadata so the LLM has full
context without relying on the frontend to pass segment_context.
"""
return (
f"Document: {analysis.doc_name}\n"
f"Standard: {analysis.standard_name}\n"
f"Finding [{finding.seq + 1}]: {finding.title}\n"
f"Status: {finding.status}\n"
f"Clause reference: {finding.clause_ref or 'N/A'}\n"
f"Description: {finding.description}\n"
f"Overall conclusion: {analysis.conclusion}\n"
)
def generate_suggestions(
finding: "FindingRecord",
analysis: "AnalysisRecord",
client: "BaseLLMClient",
) -> list[str]:
"""Generate 3 context-aware follow-up questions for a finding chat thread.
Returns exactly 3 question strings. Falls back to static templates on error.
"""
fallback = _SUGGESTION_FALLBACK.get(finding.status, _SUGGESTION_FALLBACK["warn"])
context = build_finding_context(finding, analysis)
focus = _SUGGESTION_FOCUS.get(finding.status, _SUGGESTION_FOCUS["warn"])
prompt = (
f"{context}\n\n"
f"Task: {focus}\n"
"Generate exactly 3 concise follow-up questions a compliance analyst would ask.\n"
'Return JSON: {"questions": ["question 1", "question 2", "question 3"]}\n'
"Return ONLY the JSON object."
)
response = client.chat([{"role": "user", "content": prompt}], max_tokens=300)
if not response.is_success:
return fallback
try:
result = _extract_json(response.content)
questions = result.get("questions", [])
if isinstance(questions, list) and len(questions) >= 3:
return [str(q) for q in questions[:3]]
except (ValueError, TypeError) as exc:
logger.warning("generate_suggestions JSON parse failed: {}", exc)
return fallback
def detect_cross_clause_conflicts(
findings: list[dict],
client: "BaseLLMClient",
) -> list[dict]:
"""Detect contradictions and missing cross-references across all findings.
Runs a single LLM call after all per-clause findings are collected.
Returns a list of conflict dicts: {type, finding_a, finding_b, desc}.
Returns an empty list on LLM failure so the caller can proceed without it.
"""
if len(findings) < 2:
# Need at least 2 findings to compare
return []
findings_text = "\n".join(
f"[{i+1}] [{f['status'].upper()}] {f['title']}: {f['desc']}"
+ (f" (Ref: {f['clause_ref']})" if f.get("clause_ref") else "")
for i, f in enumerate(findings)
)
prompt = (
"You are a compliance expert. Review the following compliance findings from the same document "
"and identify any cross-clause issues:\n\n"
f"Findings:\n{findings_text}\n\n"
"Return JSON array of conflicts (empty array [] if none found):\n"
"[\n"
" {\n"
' "type": "contradiction" | "missing_ref" | "cumulative_risk",\n'
' "finding_a": <1-based index>,\n'
' "finding_b": <1-based index or null>,\n'
' "desc": "Brief description of the cross-clause issue (max 100 chars)"\n'
" }\n"
"]\n"
"Return ONLY the JSON array."
)
try:
response = client.chat([{"role": "user", "content": prompt}], max_tokens=600)
if not response.is_success:
return []
result = _extract_json(response.content)
if isinstance(result, list):
return [
{
"type": str(c.get("type", "contradiction")),
"finding_a": int(c.get("finding_a", 0)),
"finding_b": c.get("finding_b"),
"desc": str(c.get("desc", "")),
}
for c in result
if isinstance(c, dict)
]
except Exception as exc:
logger.warning("detect_cross_clause_conflicts failed: {}", exc)
return []
+266 -113
View File
@@ -277,7 +277,6 @@ class DocumentCommandService:
message="Document record created",
)
temp_path = ""
try:
self.binary_store.save(
object_name=object_name,
@@ -297,117 +296,20 @@ class DocumentCommandService:
stage="store",
message="Source file stored",
)
suffix = os.path.splitext(file_name)[1]
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
temp_file.write(content)
temp_path = temp_file.name
parsed_document = self.parser.parse(
file_path=temp_path,
# Delegate parse → embed → index to the shared processing method.
# This same method is invoked by the Celery worker for async processing.
return self._process_document(
doc_id=doc_id,
doc_name=final_doc_name,
)
self._safe_mark_run_parsed(doc_id=doc_id, run_id=run_id, parsed_document=parsed_document)
artifact_keys: dict[str, str] = {}
try:
artifact_keys = self._save_parse_artifacts(doc_id=doc_id, parsed_document=parsed_document)
except Exception:
logger.warning("Parse artifact binary persistence failed for doc_id={}", doc_id)
self.document_repository.update_status(
doc_id,
DocumentStatus.PARSED,
parser_name=parsed_document.parser_name,
metadata={
"parser_backend": parsed_document.parser_name,
"parse_task_id": parsed_document.metadata.get("task_id", ""),
"layout_count": parsed_document.metadata.get("layout_count", len(parsed_document.raw_layouts)),
"structure_node_count": len(parsed_document.structure_nodes),
"semantic_block_count": len(parsed_document.semantic_blocks),
"vector_chunk_count": len(parsed_document.vector_chunks),
"artifact_keys": artifact_keys,
"processing_stage": "parsed",
},
)
current_status = DocumentStatus.PARSED
current_stage = "embed"
self._safe_replace_processing_artifacts(doc_id=doc_id, run_id=run_id, artifact_keys=artifact_keys)
self._safe_append_status_event(
doc_id=doc_id,
run_id=run_id,
from_status=DocumentStatus.STORED.value,
to_status=DocumentStatus.PARSED.value,
stage="parse",
message="Document parsed",
metadata={"artifact_count": len(artifact_keys)},
)
if self.parse_artifact_store:
try:
self.parse_artifact_store.save(
doc_id,
parsed_document.structure_nodes,
parsed_document.semantic_blocks,
)
except Exception:
logger.warning("ParseArtifactStore.save failed for doc_id={}", doc_id)
chunks = self.chunk_builder.build(
parsed_document=parsed_document,
file_name=file_name,
final_doc_name=final_doc_name,
content=content,
regulation_type=regulation_type,
version=version,
)
if not chunks:
raise ValueError("解析完成但没有生成可入库的 chunks")
vectors = self.embedding_provider.embed_texts([chunk.embedding_text for chunk in chunks])
current_stage = "index"
inserted = self.vector_index.upsert(chunks, vectors)
if inserted != len(chunks):
logger.warning("Milvus upsert count mismatched: inserted={}, chunks={}", inserted, len(chunks))
health = self.vector_index.health()
self.document_repository.update_status(
doc_id,
DocumentStatus.INDEXED,
chunk_count=len(chunks),
summary="",
summary_latency_ms=0,
index_name=health.get("collection_name", ""),
metadata={
"index_collection": health.get("collection_name", ""),
"processing_stage": "indexed",
},
)
current_status = DocumentStatus.INDEXED
index_name = health.get("collection_name", "")
self._safe_mark_run_indexed(
doc_id=doc_id,
generate_summary=generate_summary,
run_id=run_id,
chunk_count=len(chunks),
index_name=index_name,
)
self._safe_append_status_event(
doc_id=doc_id,
run_id=run_id,
from_status=DocumentStatus.PARSED.value,
to_status=DocumentStatus.INDEXED.value,
stage="index",
message="Document indexed",
metadata={"chunk_count": len(chunks), "index_name": index_name},
)
stored = self.document_repository.get(doc_id)
return DocumentProcessResult(
doc_id=doc_id,
doc_name=final_doc_name,
status=(stored.status.value if stored else DocumentStatus.INDEXED.value),
message="处理成功",
num_chunks=len(chunks),
summary=stored.summary if stored else "",
summary_latency_ms=stored.summary_latency_ms if stored else 0,
)
except Exception as exc:
logger.exception("文档处理失败: doc_id={}", doc_id)
logger.exception("文档存储失败: doc_id={}", doc_id)
failure_stage = current_stage
self.document_repository.update_status(
doc_id,
@@ -439,6 +341,183 @@ class DocumentCommandService:
status=DocumentStatus.FAILED.value,
message=f"文档处理失败: {exc}",
)
def store_document(
self,
*,
doc_id: str | None = None,
file_name: str,
content: bytes,
content_type: str,
doc_name: str | None,
regulation_type: str,
version: str,
generate_summary: bool,
) -> tuple[str, str | None]:
"""Store the binary file and create the Document record.
Returns (doc_id, run_id). Does NOT parse, embed, or index.
This is the fast synchronous first step; processing is enqueued separately.
The caller is responsible for enqueuing the follow-up process_document_task.
"""
doc_id = doc_id or str(uuid.uuid4())[:8]
final_doc_name = doc_name or file_name
object_name = f"{doc_id}/{file_name}"
document = Document(
doc_id=doc_id,
doc_name=final_doc_name,
file_name=file_name,
object_name=object_name,
content_type=content_type,
size_bytes=len(content),
regulation_type=regulation_type,
version=version,
metadata={"generate_summary": generate_summary},
)
self.document_repository.create(document)
run_id = self._safe_create_processing_run(
doc_id=doc_id, trigger_type="upload", generate_summary=generate_summary
)
self.binary_store.save(
object_name=object_name, data=content,
content_type=content_type, metadata={"doc_id": doc_id},
)
self.document_repository.update_status(doc_id, DocumentStatus.STORED)
self._safe_mark_run_stored(doc_id=doc_id, run_id=run_id)
self._safe_append_status_event(
doc_id=doc_id, run_id=run_id,
from_status=DocumentStatus.PENDING.value, to_status=DocumentStatus.STORED.value,
stage="store", message="Source file stored",
)
return doc_id, run_id
def _process_document(
self,
*,
doc_id: str,
file_name: str,
final_doc_name: str,
content: bytes,
regulation_type: str,
version: str,
generate_summary: bool,
run_id: str | None = None,
) -> DocumentProcessResult:
"""Run parse → chunk → embed → index for a document that is already stored.
Called both synchronously (from upload_and_process) and asynchronously
(from the Celery process_document_task worker). All side-effects write
through DocumentProcessingStore so callers can poll progress.
"""
current_status = DocumentStatus.STORED
current_stage = "parse"
temp_path = ""
try:
suffix = os.path.splitext(file_name)[1]
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
temp_file.write(content)
temp_path = temp_file.name
parsed_document = self.parser.parse(
file_path=temp_path,
doc_id=doc_id,
doc_name=final_doc_name,
)
self._safe_mark_run_parsed(doc_id=doc_id, run_id=run_id, parsed_document=parsed_document)
artifact_keys: dict[str, str] = {}
try:
artifact_keys = self._save_parse_artifacts(doc_id=doc_id, parsed_document=parsed_document)
except Exception:
logger.warning("Parse artifact binary persistence failed for doc_id={}", doc_id)
self.document_repository.update_status(
doc_id,
DocumentStatus.PARSED,
parser_name=parsed_document.parser_name,
metadata={
"parser_backend": parsed_document.parser_name,
"parse_task_id": parsed_document.metadata.get("task_id", ""),
"layout_count": parsed_document.metadata.get("layout_count", len(parsed_document.raw_layouts)),
"structure_node_count": len(parsed_document.structure_nodes),
"semantic_block_count": len(parsed_document.semantic_blocks),
"vector_chunk_count": len(parsed_document.vector_chunks),
"artifact_keys": artifact_keys,
"processing_stage": "parsed",
},
)
current_status = DocumentStatus.PARSED
current_stage = "embed"
self._safe_replace_processing_artifacts(doc_id=doc_id, run_id=run_id, artifact_keys=artifact_keys)
self._safe_append_status_event(
doc_id=doc_id, run_id=run_id,
from_status=DocumentStatus.STORED.value, to_status=DocumentStatus.PARSED.value,
stage="parse", message="Document parsed", metadata={"artifact_count": len(artifact_keys)},
)
if self.parse_artifact_store:
try:
self.parse_artifact_store.save(
doc_id, parsed_document.structure_nodes, parsed_document.semantic_blocks,
)
except Exception:
logger.warning("ParseArtifactStore.save failed for doc_id={}", doc_id)
chunks = self.chunk_builder.build(
parsed_document=parsed_document,
regulation_type=regulation_type,
version=version,
)
if not chunks:
raise ValueError("解析完成但没有生成可入库的 chunks")
vectors = self.embedding_provider.embed_texts([chunk.embedding_text for chunk in chunks])
current_stage = "index"
inserted = self.vector_index.upsert(chunks, vectors)
if inserted != len(chunks):
logger.warning("Milvus upsert count mismatched: inserted={}, chunks={}", inserted, len(chunks))
health = self.vector_index.health()
index_name = health.get("collection_name", "")
self.document_repository.update_status(
doc_id, DocumentStatus.INDEXED,
chunk_count=len(chunks), summary="", summary_latency_ms=0,
index_name=index_name,
metadata={"index_collection": index_name, "processing_stage": "indexed"},
)
self._safe_mark_run_indexed(doc_id=doc_id, run_id=run_id, chunk_count=len(chunks), index_name=index_name)
self._safe_append_status_event(
doc_id=doc_id, run_id=run_id,
from_status=DocumentStatus.PARSED.value, to_status=DocumentStatus.INDEXED.value,
stage="index", message="Document indexed",
metadata={"chunk_count": len(chunks), "index_name": index_name},
)
stored = self.document_repository.get(doc_id)
return DocumentProcessResult(
doc_id=doc_id, doc_name=final_doc_name,
status=(stored.status.value if stored else DocumentStatus.INDEXED.value),
message="处理成功", num_chunks=len(chunks),
summary=stored.summary if stored else "",
summary_latency_ms=stored.summary_latency_ms if stored else 0,
)
except Exception as exc:
logger.exception("文档处理失败: doc_id={}", doc_id)
self.document_repository.update_status(
doc_id, DocumentStatus.FAILED, error_message=str(exc),
metadata={"failure_reason": str(exc), "processing_stage": "failed", "failure_stage": current_stage},
)
self._safe_mark_run_failed(
doc_id=doc_id, run_id=run_id, failure_stage=current_stage, error_message=str(exc)
)
self._safe_append_status_event(
doc_id=doc_id, run_id=run_id,
from_status=current_status.value, to_status=DocumentStatus.FAILED.value,
stage=current_stage, message=str(exc),
)
return DocumentProcessResult(
doc_id=doc_id, doc_name=final_doc_name,
status=DocumentStatus.FAILED.value, message=f"文档处理失败: {exc}",
)
finally:
if temp_path and os.path.exists(temp_path):
try:
@@ -446,12 +525,29 @@ class DocumentCommandService:
except OSError:
logger.warning("临时文件清理失败: {}", temp_path)
def delete(self, doc_id: str) -> bool:
"""Delete document record, binary file, and vector chunks."""
"""Delete document record, binary file, and vector chunks.
Handles two cases:
- Normal docs: have a metadata record in the document repository.
- Milvus-only (synthetic) docs: visible in management-list because they
have Milvus vectors but no JSON/PG metadata record. We still clean up
the Milvus chunks so the document disappears from the list.
"""
document = self.document_repository.get(doc_id)
if not document:
# No metadata record — might be a Milvus-only synthetic document.
# Attempt vector cleanup directly; treat as success if any chunks deleted.
try:
deleted_count = self.vector_index.delete_by_document(doc_id)
if deleted_count > 0:
logger.info("Deleted Milvus-only doc (no metadata record): doc_id={} chunks={}", doc_id, deleted_count)
return True
except Exception as exc:
logger.warning("Milvus-only delete failed for doc_id={}: {}", doc_id, exc)
return False
# Normal doc: clean up binary, vectors, artifacts, processing records, metadata.
try:
self.binary_store.delete(document.object_name)
except Exception:
@@ -549,13 +645,16 @@ class DocumentQueryService:
result.append(doc)
# Surface Milvus-only docs that have no metadata record at all.
# MinIO almost certainly has their binaries (they were uploaded), so
# set object_name to the sentinel "{doc_id}/" so the route marks
# has_file=True; the download endpoint will list MinIO to find the file.
for doc_id, row in milvus_by_id.items():
if doc_id not in meta_by_id:
synthetic = Document(
doc_id=doc_id,
doc_name=row.get("doc_title", doc_id),
file_name=row.get("doc_title", doc_id),
object_name="",
object_name=f"{doc_id}/", # sentinel: MinIO prefix exists
content_type="",
size_bytes=0,
status=DocumentStatus.INDEXED,
@@ -568,9 +667,63 @@ class DocumentQueryService:
result.sort(key=lambda d: d.updated_at, reverse=True)
return result[:limit] if limit is not None else result
def download(self, doc_id: str) -> tuple[Document, bytes]:
"""Handle download for the Document Query Service instance."""
def download(self, doc_id: str) -> tuple["Document", bytes]:
"""Return the document record and its binary content from MinIO.
Fallback strategy for Milvus-only docs (no JSON/PG metadata record):
1. Try metadata repository first (normal path).
2. If metadata is missing, list MinIO objects with prefix ``{doc_id}/``
and synthesise a minimal Document from the first object found.
This handles documents whose metadata records were lost but whose
binary files are still in object storage.
3. If neither source has the file, raise FileNotFoundError.
"""
from app.domain.documents import Document, DocumentStatus
document = self.document_repository.get(doc_id)
if not document:
raise FileNotFoundError(f"文档不存在: {doc_id}")
if document and document.object_name and not document.object_name.endswith("/"):
# Normal doc with a concrete object_name — read directly.
return document, self.binary_store.read(document.object_name)
if document and not document.object_name:
raise FileNotFoundError(f"该文档无原始文件(仅含索引数据,无法下载): {doc_id}")
if not document or document.object_name.endswith("/"):
# Metadata missing — try to find the file in MinIO by doc_id prefix.
try:
objects = self.binary_store.list_objects(prefix=f"{doc_id}/")
# Filter out artifact JSON files; prefer the source document.
candidates = [o for o in objects if not o.endswith(".json")]
if not candidates:
candidates = objects # fall back to all objects if only JSON found
if not candidates:
raise FileNotFoundError(f"文档不存在(MinIO 和元数据均无记录): {doc_id}")
object_name = candidates[0]
file_name = object_name.split("/", 1)[-1] if "/" in object_name else object_name
# Guess content type from extension.
ext = file_name.rsplit(".", 1)[-1].lower() if "." in file_name else ""
_ct_map = {
"pdf": "application/pdf",
"docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
"doc": "application/msword",
"txt": "text/plain",
}
content_type = _ct_map.get(ext, "application/octet-stream")
# Synthesise a minimal Document so the route can build the response.
document = Document(
doc_id=doc_id,
doc_name=file_name,
file_name=file_name,
object_name=object_name,
content_type=content_type,
size_bytes=0,
status=DocumentStatus.INDEXED,
)
logger.info("MinIO fallback download: doc_id={} object={}", doc_id, object_name)
except FileNotFoundError:
raise
except Exception as exc:
raise FileNotFoundError(f"文档不存在: {doc_id}") from exc
return document, self.binary_store.read(document.object_name)
@@ -0,0 +1,255 @@
"""Orchestrates regulatory source crawlers and LLM enrichment pipeline."""
from __future__ import annotations
import hashlib
from typing import Any, Generator
from loguru import logger
from app.config.settings import settings
from app.domain.documents import ParsedDocument
from app.infrastructure.perception.base_event_store import BaseEventStore
from app.infrastructure.perception.base_notification_store import BaseNotificationStore
from app.infrastructure.perception.crawlers.base import BaseCrawler, RawEvent
from app.infrastructure.perception.llm_pipeline import LlmPipeline
from app.infrastructure.parser.local_chunk_builder import LocalRegulationChunkBuilder
def _event_id(source: str, standard_code: str) -> str:
"""Deterministic 12-char ID from source + standard_code."""
return hashlib.sha256(f"{source}-{standard_code}".encode()).hexdigest()[:12]
def _content_hash(raw_text: str) -> str:
return hashlib.sha256(raw_text.encode()).hexdigest()
def _is_significant(changed_sections: list[dict]) -> bool:
"""Report whether any changed section is worth notifying every user about.
changed_sections legitimately includes cosmetic edits — the differ
(subproject 1) still reports a fixed typo or a dropped trailing period as
a change, it just doesn't send those to the LLM. Broadcasting a
notification for every cosmetic edit would train people to ignore it, so
this reuses the same significance test the differ's own LLM gate applies:
a numeric or deontic change, or a whole paragraph added or removed.
"""
return any(
section.get("numeric_changed")
or section.get("deontic_changed")
or section.get("change_type") in ("added", "removed")
for section in changed_sections
)
def _index_in_knowledge_base(event: dict, *, embedding_provider: Any, vector_index: Any) -> None:
"""Chunk, embed, and upsert a regulation's text into the shared knowledge base.
Always uses the local markdown chunker, never get_chunk_builder() — that
bootstrap function resolves to AliyunVectorChunkBuilder when
settings.chunk_backend == "aliyun" (the deployed value), which consumes
Aliyun DocMind's structured parse output. Crawled text has no such parse
output; it is already plain text (trafilatura, subproject 1), which is
exactly what LocalRegulationChunkBuilder chunks directly.
delete_by_document runs unconditionally before upsert — a no-op for a
brand-new event, and the only way to keep a changed regulation from
leaving its superseded text retrievable alongside the new version.
"""
vector_index.delete_by_document(event["id"])
parsed = ParsedDocument(
doc_id=event["id"],
doc_name=event.get("title", ""),
structure_nodes=[],
semantic_blocks=[],
vector_chunks=[],
parser_name="perception_crawl",
raw_text=event.get("raw_text") or "",
)
builder = LocalRegulationChunkBuilder(
chunk_size=settings.chunk_size, chunk_overlap=settings.chunk_overlap,
)
chunks = builder.build(
parsed_document=parsed,
# regulation_type/version fill the same slots a manually uploaded
# document's form fields would, so the two intake paths are
# indistinguishable to retrieval and compliance analysis.
regulation_type=event.get("category", ""),
version=event.get("standard_code", ""),
)
if not chunks:
return
vectors = embedding_provider.embed_texts([c.embedding_text for c in chunks])
vector_index.upsert(chunks, vectors)
def _raw_to_dict(raw: RawEvent, event_id: str, content_hash: str, raw_text: str) -> dict:
return {
"id": event_id,
"source": raw.source,
"source_label": raw.source_label,
"standard_code": raw.standard_code,
"title": raw.title,
"summary": raw.summary,
"full_text_url": raw.full_text_url,
"status": raw.status,
"impact_level": "medium",
"published_at": raw.published_at,
"effective_at": raw.effective_at,
"category": raw.category,
"tags": raw.tags,
# Persisted so the next crawl has a baseline to diff against. Without
# this the change detector has nothing to compare and every update
# looks like a first sighting.
"raw_text": raw_text,
"content_hash": content_hash,
"previous_hash": None,
}
class CrawlService:
"""Orchestrate crawlers, hash-based change detection, and LLM enrichment."""
def __init__(
self,
crawlers: dict[str, BaseCrawler],
event_store: BaseEventStore,
llm_pipeline: LlmPipeline,
retrieval_service: Any,
notification_store: BaseNotificationStore,
embedding_provider: Any,
vector_index: Any,
) -> None:
self._crawlers = crawlers
self._store = event_store
self._pipeline = llm_pipeline
self._retrieval = retrieval_service
self._notifications = notification_store
self._embedding_provider = embedding_provider
self._vector_index = vector_index
def run_crawl(
self, sources: list[str] | None = None
) -> Generator[dict, None, None]:
"""Run crawl for selected sources. Yields SSE-ready progress dicts."""
targets = sources or list(self._crawlers.keys())
total_new = 0
total_updated = 0
for source_key in targets:
crawler = self._crawlers.get(source_key)
if not crawler:
yield {"event": "error", "data": f"Unknown source: {source_key}"}
continue
yield {"event": "progress", "data": {"source": source_key, "stage": "fetching"}}
try:
raw_events = crawler.fetch(limit=settings.perception_max_events_per_source)
except Exception as exc:
logger.exception("Crawler failed source={}", source_key)
yield {"event": "error", "data": {"source": source_key, "message": str(exc)}}
continue
yield {
"event": "progress",
"data": {"source": source_key, "stage": "processing", "fetched": len(raw_events)},
}
new_count = 0
updated_count = 0
for raw in raw_events:
eid = _event_id(raw.source, raw.standard_code)
# List pages carry only a code and a title, which is not enough
# to detect a change in the regulation itself. Fetch the body,
# degrading to whatever the list page gave us if that fails.
body_text = crawler.fetch_full_text(raw.full_text_url) or raw.raw_text or raw.title
new_hash = _content_hash(body_text)
existing = self._store.get(eid)
if existing and existing.get("content_hash") == new_hash:
continue
is_update = existing is not None
old_body = existing.get("raw_text") or "" if is_update else ""
previous_hash = existing.get("content_hash") if is_update else None
event_dict = _raw_to_dict(raw, eid, new_hash, body_text)
event_dict["previous_hash"] = previous_hash
try:
structure = self._pipeline.extract_structure(event_dict)
event_dict.update(structure)
except Exception as exc:
logger.warning("Structure extraction failed id={} err={}", eid, exc)
try:
affected = self._pipeline.assess_impact(event_dict, self._retrieval)
event_dict["affected_docs"] = affected
except Exception as exc:
logger.warning("Impact assessment failed id={} err={}", eid, exc)
# Events stored before raw_text was persisted have no baseline,
# so they are treated as a first sighting and establish one now.
if is_update and old_body and body_text:
try:
diff = self._pipeline.compute_diff(old_body, body_text)
event_dict["change_summary"] = diff.get("change_summary")
event_dict["changed_sections"] = diff.get("changed_sections")
except Exception as exc:
logger.warning("Diff failed id={} err={}", eid, exc)
self._store.upsert(event_dict)
should_index = not is_update or _is_significant(event_dict.get("changed_sections") or [])
try:
if not is_update:
self._notifications.create(
event_id=eid, kind="new", title=raw.title,
impact_level=event_dict.get("impact_level"), summary=None,
)
elif should_index: # significant change, already computed above
self._notifications.create(
event_id=eid, kind="changed", title=raw.title,
impact_level=event_dict.get("impact_level"),
summary=event_dict.get("change_summary"),
)
except Exception as exc:
logger.warning("Notification create failed id={} err={}", eid, exc)
if should_index:
try:
_index_in_knowledge_base(
event_dict,
embedding_provider=self._embedding_provider,
vector_index=self._vector_index,
)
except Exception as exc:
logger.warning("Knowledge base indexing failed id={} err={}", eid, exc)
if is_update:
updated_count += 1
else:
new_count += 1
total_new += new_count
total_updated += updated_count
yield {
"event": "progress",
"data": {
"source": source_key,
"stage": "done",
"new": new_count,
"updated": updated_count,
},
}
yield {
"event": "done",
"data": {"total_new": total_new, "total_updated": total_updated},
}
@@ -6,7 +6,7 @@ import json
from typing import Generator
from app.application.knowledge.services import KnowledgeRetrievalService
from app.infrastructure.perception.mock_event_store import MockEventStore
from app.infrastructure.perception.base_event_store import BaseEventStore
from app.services.llm.llm_factory import get_llm_client
from app.config.settings import settings
@@ -22,7 +22,7 @@ class PerceptionService:
def __init__(
self,
event_store: MockEventStore,
event_store: BaseEventStore,
retrieval_service: KnowledgeRetrievalService,
) -> None:
self._store = event_store
+112 -1
View File
@@ -82,6 +82,34 @@ class Settings(BaseSettings):
parser_backend: str = Field(default="aliyun", description="解析后端(local/aliyun)")
chunk_backend: str = Field(default="aliyun", description="分块后端(local/aliyun)")
document_repository_backend: str = Field(default="json", description="文档元数据存储后端 (json/postgres)")
# When True, document processing is enqueued to Celery workers via Redis.
# When False (default), processing runs in a FastAPI BackgroundTask in the same process —
# no external worker needed. Switch to True only when a Celery worker is running.
use_celery_worker: bool = Field(default=False, description="使用 Celery Worker 异步处理文档 (需要 Worker 运行中)")
# ── Perception crawl ──────────────────────────────────────────────────────
perception_crawl_timeout_seconds: int = Field(
default=120, description="HTTP timeout for regulatory source crawlers."
)
perception_max_events_per_source: int = Field(
default=100, description="Maximum events fetched per source per crawl run."
)
perception_diff_min_change_ratio: float = Field(
default=0.02,
description=(
"Fraction of characters that must differ before an otherwise "
"unremarkable paragraph edit is worth an LLM classification call. "
"Numeric and deontic changes bypass this gate entirely."
),
)
perception_crawl_interval_seconds: int = Field(
default=21600,
description=(
"How often Celery Beat runs the scheduled crawl-all-sources task, "
"in seconds. Default 21600 = 6 hours. Only takes effect when a "
"Beat process is running (./dev.sh start beat)."
),
)
# Keep configuration setup explicit so runtime behavior is easy to reason about.
api_host: str = Field(default="0.0.0.0", description="API服务地址")
@@ -101,7 +129,7 @@ class Settings(BaseSettings):
# Keep configuration setup explicit so runtime behavior is easy to reason about.
qwen_api_key: str = Field(default="", description="Qwen API密钥")
qwen_base_url: str = Field(default="http://6.86.80.4:30080/v1", description="Qwen API地址")
qwen_model: str = Field(default="qwen3.5-flash", description="Qwen文本模型")
qwen_model: str = Field(default="qwen3.6-flash", description="Qwen文本模型")
qwen_vl_model: str = Field(default="qwen3-vl-plus", description="Qwen视觉模型")
# Keep configuration setup explicit so runtime behavior is easy to reason about.
@@ -109,6 +137,7 @@ class Settings(BaseSettings):
rag_retrieval_top_k: int = Field(default=20, description="精排前召回候选数量(reranker 启用时生效)")
rag_max_context_tokens: int = Field(default=2000, description="RAG最大上下文token数")
rag_summary_max_tokens: int = Field(default=10240, description="文档摘要最大token数")
rag_skills_max_tokens: int = Field(default=2048, description="技能类 RAG 最大 token 数")
reranker_enabled: bool = Field(default=False, description="是否启用 Cross-Encoder 精排")
reranker_base_url: str = Field(default="", description="Reranker API 地址")
@@ -116,6 +145,42 @@ class Settings(BaseSettings):
reranker_api_key: str = Field(default="", description="Reranker API 密钥")
reranker_top_k: int = Field(default=5, description="精排后保留的最终结果数量")
# ── HyDE (Hypothetical Document Embeddings) ──────────────────────────────
# When enabled, the agentic and standard RAG pipelines generate a short
# hypothetical answer before retrieval, then embed that text instead of the
# raw query. This closes the vocabulary gap between terse queries and longer
# document passages, typically improving recall by 15-30% on vague queries.
hyde_enabled: bool = Field(default=True, description="启用 HyDE 查询增强(假设文档嵌入)")
hyde_max_tokens: int = Field(default=200, description="HyDE 假设段落最大 token 数")
# Use a lightweight model for HyDE to reduce latency and cost.
# HyDE only needs a short plausible passage — a fast cheap model is sufficient.
# Leave empty to fall back to the main llm_provider / llm_model.
hyde_llm_provider: str = Field(default="", description="HyDE 专用 LLM 提供商(空则复用主 LLM)")
hyde_llm_model: str = Field(default="", description="HyDE 专用 LLM 模型(空则复用主 LLM)")
# ── Agentic RAG (P0-1) ───────────────────────────────────────────────────
# Controls the multi-step reasoning pipeline exposed at /agent/agentic/stream.
agentic_max_sub_queries: int = Field(
default=4,
description="Agentic 模式最大子查询分解数量(compare / multi_hop 意图触发)",
)
agentic_grounding_threshold: float = Field(
default=0.65,
description=(
"引文锚定 fast-path 阈值:avg_score > 此值且 chunks ≥ 3 时跳过 LLM grounding check"
"直接判定为充分;降低此值可让更多问题触发 LLM 二次验证。"
),
)
agentic_intent_max_tokens: int = Field(
default=200, description="意图分析步骤 LLM 最大 token 数"
)
agentic_plan_max_tokens: int = Field(
default=400, description="查询分解步骤 LLM 最大 token 数"
)
agentic_grounding_max_tokens: int = Field(
default=250, description="引文锚定步骤 LLM 最大 token 数"
)
# Keep configuration setup explicit so runtime behavior is easy to reason about.
milvus_index_type: str = Field(default="IVF_FLAT", description="Milvus索引类型")
milvus_nlist: int = Field(default=128, description="Milvus nlist参数")
@@ -124,6 +189,52 @@ class Settings(BaseSettings):
# Keep configuration setup explicit so runtime behavior is easy to reason about.
session_max_sessions: int = Field(default=100, description="最大会话数量")
session_timeout_minutes: int = Field(default=30, description="会话超时时间(分钟)")
session_backend: str = Field(
default="memory",
description="会话存储后端 (memory | redis)。redis 需要 Redis 可用。",
)
# ── Auth ──────────────────────────────────────────────────────────────────
# Generate a strong secret: python -c "import secrets; print(secrets.token_hex(32))"
auth_secret_key: str = Field(
default="change-me-in-production-must-be-32-or-more-characters-long",
description="JWT signing secret. MUST be changed in production.",
)
auth_algorithm: str = Field(default="HS256", description="JWT signing algorithm.")
auth_token_expire_minutes: int = Field(default=480, description="JWT TTL in minutes (default 8 hours).")
auth_enabled: bool = Field(default=True, description="Set False to bypass auth (development only).")
# ── CORS ──────────────────────────────────────────────────────────────────
cors_allow_origins: str = Field(
default="http://localhost:5173",
description="Comma-separated allowed CORS origins. Never use * in production.",
)
# ── MCP ───────────────────────────────────────────────────────────────────
# The MCP SDK enables DNS-rebinding protection whenever the transport is
# bound to a loopback host, which rejects any Host header not in this list
# with HTTP 421. Deployments reachable by a real hostname/IP must list it
# here or every remote MCP client is refused before the handler runs.
mcp_allowed_hosts: str = Field(
default="127.0.0.1:*,localhost:*,[::1]:*",
description=(
"Comma-separated Host header values accepted by the MCP endpoint. "
"A ':*' suffix matches any port. Set to '*' to disable DNS-rebinding "
"protection entirely (not recommended)."
),
)
# Optional override for the URL shown on the System Status page and copied
# into client configs. Needed because request.base_url reflects the Host
# header, which the Vite dev proxy (changeOrigin: true) and reverse proxies
# that do not forward the original Host both rewrite.
mcp_public_url: str = Field(
default="",
description=(
"Externally reachable MCP endpoint URL, e.g. http://6.86.80.9:8000/mcp/. "
"Leave empty to derive it from the incoming request."
),
)
@lru_cache
def get_settings() -> Settings:
+10
View File
@@ -0,0 +1,10 @@
"""Auth domain: role definitions and token claim models.
The domain layer defines what a user identity looks like (UserClaims) and
what roles exist (UserRole). Infrastructure details (JWT, bcrypt, PostgreSQL)
live under infrastructure/auth and never leak into this package.
"""
from .models import UserClaims, UserRole
__all__ = ["UserClaims", "UserRole"]
+42
View File
@@ -0,0 +1,42 @@
"""Auth domain models: roles and token claims.
UserRole defines the four roles from PPT Slide 12.
UserClaims is what the JWT decodes to — it is the identity object passed
through FastAPI dependency injection to route handlers.
"""
from __future__ import annotations
import enum
from dataclasses import dataclass
class UserRole(str, enum.Enum):
"""Access roles mirroring the four-role RBAC matrix from the product spec.
ADMIN — full platform access including system management.
LEGAL — knowledge query, document review, compliance checks.
EHS — knowledge query, perception/regulatory signals.
READONLY — knowledge query only.
"""
ADMIN = "admin"
LEGAL = "legal"
EHS = "ehs"
READONLY = "readonly"
@dataclass
class UserClaims:
"""Decoded JWT payload representing an authenticated user.
Instances are created by JWTHandler.decode_token() and injected into
route handlers via the get_current_user FastAPI dependency.
"""
# Unique user identifier (UUID string stored in PostgreSQL users table).
user_id: str
# Display name used for audit log entries.
username: str
# Role determines which resources the user may access.
role: UserRole
+66
View File
@@ -0,0 +1,66 @@
"""Domain ports for compliance history persistence."""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional
@dataclass
class FindingRecord:
"""Single finding row linked to an analysis."""
id: str
analysis_id: str
seq: int
title: str
description: str
status: str # "ok" | "warn" | "risk"
clause_ref: Optional[str] = None
@dataclass
class AnalysisRecord:
"""Full compliance analysis record with nested findings."""
id: str # UUID string; empty string means not yet persisted
created_at: datetime
created_by: Optional[str]
doc_name: str
standard_name: str
risk_score: int
conclusion: str
actions: list # list[dict] — serialised action items
para_text: str
highlight_terms: list # list[str]
findings: list[FindingRecord] = field(default_factory=list)
class ComplianceRepository(ABC):
"""Port for persisting and retrieving compliance analysis records."""
@abstractmethod
def save_analysis(self, record: AnalysisRecord) -> str:
"""Persist a new analysis record and return the assigned UUID string."""
@abstractmethod
def list_analyses(self, limit: int = 50, offset: int = 0) -> list[AnalysisRecord]:
"""Return analyses ordered by created_at DESC, without nested findings."""
@abstractmethod
def get_analysis(self, analysis_id: str) -> Optional[AnalysisRecord]:
"""Return a single analysis with all nested findings, or None."""
@abstractmethod
def delete_analysis(self, analysis_id: str) -> None:
"""Delete an analysis and all related findings and chat messages (cascade)."""
@abstractmethod
def save_message(self, analysis_id: str, finding_id: str, role: str, content: str) -> str:
"""Persist a chat message and return its UUID string."""
@abstractmethod
def get_messages(self, finding_id: str) -> list[dict]:
"""Return chat messages for a finding ordered by created_at ASC.
Each dict has keys: id, role, content, created_at (ISO string).
"""
@@ -0,0 +1,5 @@
"""JWT token creation and validation infrastructure.
JWTHandler is the only component in this package. It is wired through
shared/bootstrap.py and injected into FastAPI dependencies.
"""
@@ -0,0 +1,82 @@
"""JWT access token creation and decoding.
Uses python-jose for HS256 token signing. Token expiry is enforced at
decode time so expired tokens are rejected even if the signature is valid.
"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta
from typing import Any
from jose import JWTError, jwt
from loguru import logger
from app.domain.auth.models import UserClaims, UserRole
class JWTHandler:
"""Create and validate HS256 JWT access tokens.
A single shared instance is wired by bootstrap.py. Use
get_jwt_handler() from shared.bootstrap for all token operations.
"""
def __init__(
self,
*,
secret_key: str,
algorithm: str = "HS256",
expire_minutes: int = 480,
) -> None:
"""Initialise the handler with signing credentials and token lifetime."""
self._secret = secret_key
self._algorithm = algorithm
self._expire_minutes = expire_minutes
def create_access_token(
self,
*,
user_id: str,
username: str,
role: str,
) -> str:
"""Return a signed JWT containing user identity and role claims."""
now = datetime.now(UTC)
payload: dict[str, Any] = {
"sub": user_id,
"username": username,
"role": role,
"iat": now,
"exp": now + timedelta(minutes=self._expire_minutes),
}
return jwt.encode(payload, self._secret, algorithm=self._algorithm)
def decode_token(self, token: str) -> UserClaims:
"""Decode and validate a JWT, returning UserClaims.
Raises ValueError with a descriptive message on expiry, tampering,
or any other validation failure so callers do not need to know jose.
"""
try:
payload = jwt.decode(token, self._secret, algorithms=[self._algorithm])
except JWTError as exc:
msg = str(exc).lower()
if "expired" in msg:
raise ValueError("Token expired") from exc
raise ValueError(f"Invalid token: {exc}") from exc
user_id = payload.get("sub")
username = payload.get("username", "")
role_str = payload.get("role", UserRole.READONLY.value)
if not user_id:
raise ValueError("Token missing subject claim")
try:
role = UserRole(role_str)
except ValueError:
logger.warning("Unknown role in token: {}, defaulting to readonly", role_str)
role = UserRole.READONLY
return UserClaims(user_id=user_id, username=username, role=role)
@@ -0,0 +1,113 @@
"""PostgreSQL-backed user store for authentication.
Manages a `users` table with hashed passwords and roles.
Provides lookup by username for the login flow.
Table DDL is auto-applied on first connection.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
import psycopg2
import psycopg2.extras
from loguru import logger
from passlib.context import CryptContext
from app.config.settings import settings
# bcrypt context — work factor 12 is a good production default.
_PWD_CTX = CryptContext(schemes=["bcrypt"], deprecated="auto")
# DDL executed once to ensure the table exists.
_CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS users (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
username VARCHAR(100) UNIQUE NOT NULL,
hashed_pw TEXT NOT NULL,
role VARCHAR(50) NOT NULL DEFAULT 'readonly',
is_active BOOLEAN NOT NULL DEFAULT TRUE,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
"""
@dataclass
class UserRecord:
"""A single row from the users table."""
id: str
username: str
hashed_pw: str
role: str
is_active: bool
class PostgresUserStore:
"""Read and verify users stored in the PostgreSQL users table.
The connection is opened on first use and shared for the lifetime
of the singleton instance wired by bootstrap.
"""
def __init__(self) -> None:
"""Initialise the store and ensure the users table exists."""
self._conn = psycopg2.connect(
host=settings.postgres_host,
port=settings.postgres_port,
user=settings.postgres_user,
password=settings.postgres_password,
dbname=settings.postgres_db,
cursor_factory=psycopg2.extras.RealDictCursor,
)
self._conn.autocommit = True
self._ensure_table()
def _ensure_table(self) -> None:
"""Create the users table if it does not already exist."""
with self._conn.cursor() as cur:
# Enable pgcrypto so gen_random_uuid() is available for UUID primary keys.
try:
cur.execute("CREATE EXTENSION IF NOT EXISTS pgcrypto;")
except Exception:
self._conn.rollback()
cur.execute(_CREATE_TABLE_SQL)
def get_by_username(self, username: str) -> Optional[UserRecord]:
"""Return a UserRecord for the given username, or None if not found."""
with self._conn.cursor() as cur:
cur.execute(
"SELECT id, username, hashed_pw, role, is_active "
"FROM users WHERE username = %s",
(username,),
)
row = cur.fetchone()
if row is None:
return None
return UserRecord(
id=str(row["id"]),
username=row["username"],
hashed_pw=row["hashed_pw"],
role=row["role"],
is_active=row["is_active"],
)
def verify_password(self, plain: str, hashed: str) -> bool:
"""Return True if `plain` matches the stored bcrypt hash."""
return _PWD_CTX.verify(plain, hashed)
def authenticate(self, username: str, password: str) -> Optional[UserRecord]:
"""Return the UserRecord if credentials are valid, else None."""
user = self.get_by_username(username)
if user is None or not user.is_active:
return None
if not self.verify_password(password, user.hashed_pw):
return None
return user
@staticmethod
def hash_password(plain: str) -> str:
"""Hash a plain-text password with bcrypt."""
return _PWD_CTX.hash(plain)
@@ -0,0 +1,101 @@
"""DOCX report generator for compliance analysis results.
Uses python-docx (already in requirements.txt). Returns raw bytes so the
caller can stream the response without writing to disk.
"""
from __future__ import annotations
from datetime import datetime, timezone
from io import BytesIO
from docx import Document
from docx.shared import Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH
from app.domain.compliance.ports import AnalysisRecord
_STATUS_LABEL = {"ok": "Compliant", "warn": "Warning", "risk": "Non-Compliant"}
_STATUS_COLOR = {
"ok": RGBColor(0x22, 0x8B, 0x22),
"warn": RGBColor(0xFF, 0x8C, 0x00),
"risk": RGBColor(0xDC, 0x14, 0x3C),
}
def generate_docx(record: AnalysisRecord) -> bytes:
"""Generate a compliance report DOCX and return its raw bytes.
Structure:
- Cover: document name, standard, date, risk score
- Executive summary (conclusion)
- Findings table
- Recommended actions
- Footer note
"""
doc = Document()
# ── Cover ──────────────────────────────────────────────────────────────────
title_para = doc.add_heading("Compliance Analysis Report", level=0)
title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
doc.add_paragraph("")
meta_table = doc.add_table(rows=4, cols=2)
meta_table.style = "Table Grid"
labels = ["Document", "Standard", "Date", "Risk Score"]
values = [
record.doc_name,
record.standard_name,
record.created_at.strftime("%Y-%m-%d %H:%M UTC") if record.created_at else "",
f"{record.risk_score} / 100",
]
for i, (label, value) in enumerate(zip(labels, values)):
meta_table.cell(i, 0).text = label
meta_table.cell(i, 1).text = value
# ── Executive Summary ──────────────────────────────────────────────────────
doc.add_heading("Executive Summary", level=1)
doc.add_paragraph(record.conclusion)
# ── Findings ───────────────────────────────────────────────────────────────
doc.add_heading("Findings", level=1)
if record.findings:
table = doc.add_table(rows=1, cols=4)
table.style = "Table Grid"
hdr = table.rows[0].cells
for i, h in enumerate(["#", "Status", "Title", "Description / Clause"]):
hdr[i].text = h
for run in hdr[i].paragraphs[0].runs:
run.bold = True
for f in record.findings:
row = table.add_row().cells
row[0].text = str(f.seq + 1)
row[1].text = _STATUS_LABEL.get(f.status, f.status)
row[2].text = f.title
desc = f.description
if f.clause_ref:
desc += f"\n[{f.clause_ref}]"
row[3].text = desc
else:
doc.add_paragraph("No findings recorded.")
# ── Recommended Actions ────────────────────────────────────────────────────
doc.add_heading("Recommended Actions", level=1)
for i, action in enumerate(record.actions, start=1):
label = action.get("label", "Action")
value = action.get("value", "")
doc.add_paragraph(f"{i}. {label}: {value}", style="List Number")
# ── Footer note ────────────────────────────────────────────────────────────
doc.add_paragraph("")
footer = doc.add_paragraph(
f"Generated by AI Regulation Analysis System — {datetime.now(timezone.utc).strftime('%Y-%m-%d')}"
)
footer.alignment = WD_ALIGN_PARAGRAPH.CENTER
for run in footer.runs:
run.font.size = Pt(8)
run.font.color.rgb = RGBColor(0x88, 0x88, 0x88)
buf = BytesIO()
doc.save(buf)
return buf.getvalue()
@@ -0,0 +1,280 @@
# backend/app/infrastructure/compliance/repository.py
"""PostgreSQL-backed compliance analysis repository.
Follows the same psycopg2 pattern as PostgresDocumentRepository:
ThreadedConnectionPool + RealDictCursor for reads, _ensure_schema on init.
"""
from __future__ import annotations
import json
from contextlib import contextmanager
from datetime import datetime
from typing import Optional
import psycopg2
import psycopg2.extras
import psycopg2.pool
from loguru import logger
from app.domain.compliance.ports import (
AnalysisRecord,
ComplianceRepository,
FindingRecord,
)
class PostgresComplianceRepository(ComplianceRepository):
"""Stores compliance analyses, findings, and finding chat messages in PostgreSQL."""
def __init__(
self,
host: str,
port: int,
user: str,
password: str,
dbname: str,
minconn: int = 1,
maxconn: int = 5,
) -> None:
self._pool = psycopg2.pool.ThreadedConnectionPool(
minconn=minconn,
maxconn=maxconn,
host=host,
port=port,
user=user,
password=password,
dbname=dbname,
)
self._ensure_schema()
@contextmanager
def _conn(self):
conn = self._pool.getconn()
try:
yield conn
finally:
self._pool.putconn(conn)
def _ensure_schema(self) -> None:
"""Create tables if they do not exist."""
with self._conn() as conn:
with conn.cursor() as cur:
cur.execute("""
CREATE TABLE IF NOT EXISTS compliance_analyses (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
created_by VARCHAR(255),
doc_name VARCHAR(500),
standard_name VARCHAR(500),
risk_score INTEGER,
conclusion TEXT,
actions JSONB,
para_text TEXT,
highlight_terms JSONB
);
""")
cur.execute("""
CREATE TABLE IF NOT EXISTS compliance_findings (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
analysis_id UUID NOT NULL REFERENCES compliance_analyses(id) ON DELETE CASCADE,
seq INTEGER NOT NULL,
title VARCHAR(500),
description TEXT,
status VARCHAR(50),
clause_ref VARCHAR(200)
);
""")
cur.execute("""
CREATE TABLE IF NOT EXISTS finding_chat_messages (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
analysis_id UUID NOT NULL REFERENCES compliance_analyses(id) ON DELETE CASCADE,
finding_id UUID NOT NULL REFERENCES compliance_findings(id) ON DELETE CASCADE,
role VARCHAR(20) NOT NULL,
content TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
""")
conn.commit()
def save_analysis(self, record: AnalysisRecord) -> str:
"""Insert analysis + findings; return the new analysis UUID."""
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
INSERT INTO compliance_analyses
(created_by, doc_name, standard_name, risk_score,
conclusion, actions, para_text, highlight_terms)
VALUES
(%(created_by)s, %(doc_name)s, %(standard_name)s, %(risk_score)s,
%(conclusion)s, %(actions)s, %(para_text)s, %(highlight_terms)s)
RETURNING id
""",
{
"created_by": record.created_by,
"doc_name": record.doc_name,
"standard_name": record.standard_name,
"risk_score": record.risk_score,
"conclusion": record.conclusion,
"actions": json.dumps(record.actions, ensure_ascii=False),
"para_text": record.para_text,
"highlight_terms": json.dumps(record.highlight_terms, ensure_ascii=False),
},
)
row = cur.fetchone()
analysis_id = str(row["id"])
if record.findings:
with conn.cursor() as cur:
for f in record.findings:
cur.execute(
"""
INSERT INTO compliance_findings
(analysis_id, seq, title, description, status, clause_ref)
VALUES
(%(analysis_id)s, %(seq)s, %(title)s, %(desc)s, %(status)s, %(clause_ref)s)
""",
{
"analysis_id": analysis_id,
"seq": f.seq,
"title": f.title,
"desc": f.description,
"status": f.status,
"clause_ref": f.clause_ref,
},
)
conn.commit()
return analysis_id
def list_analyses(self, limit: int = 50, offset: int = 0) -> list[AnalysisRecord]:
"""Return analyses without nested findings, ordered newest first."""
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
SELECT id, created_at, created_by, doc_name, standard_name,
risk_score, conclusion, actions, para_text, highlight_terms
FROM compliance_analyses
ORDER BY created_at DESC
LIMIT %(limit)s OFFSET %(offset)s
""",
{"limit": limit, "offset": offset},
)
rows = cur.fetchall()
return [self._row_to_record(dict(r)) for r in rows]
def get_analysis(self, analysis_id: str) -> Optional[AnalysisRecord]:
"""Return analysis with nested findings list."""
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"SELECT * FROM compliance_analyses WHERE id = %(id)s",
{"id": analysis_id},
)
row = cur.fetchone()
if not row:
return None
record = self._row_to_record(dict(row))
cur.execute(
"""
SELECT id, analysis_id, seq, title, description, status, clause_ref
FROM compliance_findings
WHERE analysis_id = %(id)s
ORDER BY seq
""",
{"id": analysis_id},
)
findings = [
FindingRecord(
id=str(r["id"]),
analysis_id=str(r["analysis_id"]),
seq=r["seq"],
title=r["title"] or "",
description=r["description"] or "",
status=r["status"] or "ok",
clause_ref=r["clause_ref"],
)
for r in cur.fetchall()
]
record.findings = findings
return record
def delete_analysis(self, analysis_id: str) -> None:
"""Delete analysis; findings and chat messages cascade automatically."""
with self._conn() as conn:
with conn.cursor() as cur:
cur.execute(
"DELETE FROM compliance_analyses WHERE id = %(id)s",
{"id": analysis_id},
)
conn.commit()
def save_message(self, analysis_id: str, finding_id: str, role: str, content: str) -> str:
"""Persist a chat message; return its UUID."""
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
INSERT INTO finding_chat_messages
(analysis_id, finding_id, role, content)
VALUES
(%(analysis_id)s, %(finding_id)s, %(role)s, %(content)s)
RETURNING id
""",
{
"analysis_id": analysis_id,
"finding_id": finding_id,
"role": role,
"content": content,
},
)
row = cur.fetchone()
conn.commit()
return str(row["id"])
def get_messages(self, finding_id: str) -> list[dict]:
"""Return messages for a finding, oldest first."""
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
SELECT id, role, content, created_at
FROM finding_chat_messages
WHERE finding_id = %(finding_id)s
ORDER BY created_at ASC
""",
{"finding_id": finding_id},
)
rows = cur.fetchall()
return [
{
"id": str(r["id"]),
"role": r["role"],
"content": r["content"],
"created_at": r["created_at"].isoformat() if r["created_at"] else "",
}
for r in rows
]
def _row_to_record(self, row: dict) -> AnalysisRecord:
"""Convert a RealDictCursor row to an AnalysisRecord (no findings)."""
actions = row.get("actions") or []
if isinstance(actions, str):
actions = json.loads(actions)
highlight_terms = row.get("highlight_terms") or []
if isinstance(highlight_terms, str):
highlight_terms = json.loads(highlight_terms)
return AnalysisRecord(
id=str(row["id"]),
created_at=row["created_at"] if isinstance(row["created_at"], datetime) else datetime.utcnow(),
created_by=row.get("created_by"),
doc_name=row.get("doc_name") or "",
standard_name=row.get("standard_name") or "",
risk_score=int(row.get("risk_score") or 0),
conclusion=row.get("conclusion") or "",
actions=actions,
para_text=row.get("para_text") or "",
highlight_terms=highlight_terms,
findings=[],
)
@@ -3,11 +3,13 @@
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
@@ -45,6 +47,8 @@ class OpenAICompatibleEmbeddingProvider(EmbeddingProvider):
"""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={
@@ -56,9 +60,28 @@ class OpenAICompatibleEmbeddingProvider(EmbeddingProvider):
)
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]]:
@@ -9,10 +9,12 @@ from app.config.settings import settings
from app.domain.conversation import AnswerGenerator, AnswerResult, AnswerSource
from app.domain.retrieval import RetrievedChunk
from app.services.llm.llm_factory import get_llm_client
from app.services.rag.prompt_templates import PromptTemplates
# Keep adapter behavior explicit so integration details remain easy to audit.
PROMPT_TEMPLATES = {
# Fallback system prompts used when no rich template matches.
_FALLBACK_PROMPTS = {
"default": "你是法规知识问答助手。请仅依据提供的上下文回答;如果上下文不足,明确说明。",
"compliance_qa": "你是法规合规问答助手。优先引用给定法规原文,回答要准确、克制,并注明依据来源。",
}
@@ -38,33 +40,80 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
retrieved_chunks: list[RetrievedChunk],
history: list[dict[str, str]] | None,
prompt_template: str | None,
context_text: str | None = None,
context_filename: str | None = None,
) -> tuple[list[dict[str, str]], int]:
"""Handle build messages for this module for the Open A I Compatible Answer Generator instance."""
system_prompt = PROMPT_TEMPLATES.get(prompt_template or "compliance_qa", PROMPT_TEMPLATES["default"])
"""Build the message list to send to the LLM.
When context_text is provided the user's document is injected as a
dedicated section BEFORE the retrieved regulation chunks so the LLM
can reason about the document directly while still referencing regulations.
The retrieval step uses only the user's question, not the document text,
so embedding quality is preserved.
System prompt selection priority:
1. Rich template from PromptTemplates (compliance_qa / comparison /
compliance_check / clause_interpretation / …)
2. Fallback hardcoded prompt when no rich template matches.
"""
# Look up the rich template first; fall back to simple hardcoded prompts.
tpl_name = prompt_template or "compliance_qa"
rich_tpl = PromptTemplates.get_template(tpl_name)
if rich_tpl:
system_prompt = rich_tpl.system_prompt
else:
system_prompt = _FALLBACK_PROMPTS.get(tpl_name, _FALLBACK_PROMPTS["default"])
context_blocks = []
context_tokens = 0
# ── User document context (if attached) ───────────────────────────────
if context_text and context_text.strip():
doc_label = f"附件文档:{context_filename}" if context_filename else "附件文档"
doc_block = f"[{doc_label}]\n{context_text.strip()}"
doc_tokens = self._estimate_tokens(doc_block)
# Reserve at most half the context budget for the user document
half_budget = settings.rag_max_context_tokens // 2
if doc_tokens > half_budget:
# Truncate document to fit half the budget
ratio = half_budget / doc_tokens
doc_block = doc_block[: int(len(doc_block) * ratio)] + "\n…(文档已截断)"
doc_tokens = half_budget
context_blocks.append(doc_block)
context_tokens += doc_tokens
# ── Retrieved regulation chunks ────────────────────────────────────────
remaining_budget = settings.rag_max_context_tokens - context_tokens
for idx, chunk in enumerate(retrieved_chunks, start=1):
block = (
f"[{idx}] 文档: {chunk.doc_title}\n"
f"[法规{idx}] 文档: {chunk.doc_title}\n"
f"章节: {chunk.section_title or '未标注'}\n"
f"页码: {chunk.page_start}" + (f"-{chunk.page_end}" if chunk.page_end and chunk.page_end != chunk.page_start else "") + "\n"
f"内容: {chunk.text}"
)
block_tokens = self._estimate_tokens(block)
if context_tokens + block_tokens > settings.rag_max_context_tokens:
if block_tokens > remaining_budget:
break
remaining_budget -= block_tokens
context_tokens += block_tokens
context_blocks.append(block)
context = "\n\n".join(context_blocks)
messages = [{"role": "system", "content": system_prompt}]
for item in history or []:
messages.append({"role": item["role"], "content": item["content"]})
messages.append(
{
"role": "user",
"content": f"问题:{query}\n\n参考上下文:\n{context}\n\n请在回答后给出简要引用编号。",
}
# Craft the user turn differently when a document is attached
if context_text and context_text.strip():
user_content = (
f"问题:{query}\n\n"
f"请先基于上方附件文档内容进行分析,再结合法规参考上下文给出合规评估。"
f"\n\n参考上下文:\n{context}\n\n"
f"请在回答中注明引用来源编号(如适用)。"
)
else:
user_content = f"问题:{query}\n\n参考上下文:\n{context}\n\n请在回答后给出简要引用编号。"
messages.append({"role": "user", "content": user_content})
return messages, context_tokens
def _is_context_truncated(self, *, retrieved_chunks: list[RetrievedChunk], context_tokens: int) -> bool:
@@ -112,6 +161,8 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
provider: str | None = None,
model: str | None = None,
prompt_template: str | None = None,
context_text: str | None = None,
context_filename: str | None = None,
) -> AnswerResult:
"""Handle generate for the Open A I Compatible Answer Generator instance."""
start = time.time()
@@ -120,6 +171,8 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
retrieved_chunks=retrieved_chunks,
history=history,
prompt_template=prompt_template,
context_text=context_text,
context_filename=context_filename,
)
client = get_llm_client(provider=provider or settings.llm_provider, model=model or settings.llm_model)
response = client.chat(messages)
@@ -147,6 +200,8 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
provider: str | None = None,
model: str | None = None,
prompt_template: str | None = None,
context_text: str | None = None,
context_filename: str | None = None,
) -> Generator[dict, None, AnswerResult]:
"""Stream generate for the Open A I Compatible Answer Generator instance."""
start = time.time()
@@ -155,6 +210,8 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
retrieved_chunks=retrieved_chunks,
history=history,
prompt_template=prompt_template,
context_text=context_text,
context_filename=context_filename,
)
sources = [source.__dict__ for source in self._sources(retrieved_chunks)]
yield {"event": "sources", "data": sources}
@@ -0,0 +1,39 @@
"""Abstract base class for regulatory event stores."""
from __future__ import annotations
from abc import ABC, abstractmethod
class BaseEventStore(ABC):
"""Port interface for regulatory event persistence."""
@abstractmethod
def all(self) -> list[dict]:
"""Return all events, most-recent first."""
@abstractmethod
def get(self, event_id: str) -> dict | None:
"""Return a single event by ID, or None."""
@abstractmethod
def filter(
self,
*,
source: str | None = None,
impact_level: str | None = None,
limit: int = 50,
) -> list[dict]:
"""Return filtered events sorted by published_at descending."""
@abstractmethod
def stats(self) -> dict:
"""Return {total, high_impact, medium_impact, low_impact, recent_90d}."""
@abstractmethod
def upsert(self, event: dict) -> None:
"""Insert or update an event record."""
@abstractmethod
def get_by_standard_code(self, standard_code: str) -> dict | None:
"""Return the most-recent event with matching standard_code, or None."""
@@ -0,0 +1,47 @@
"""Abstract base class for in-app regulatory-signal notifications.
A notification is created once per triggering event (a brand-new regulation,
or a significant change to an existing one) and broadcast to every logged-in
user. There is no per-user subscription targeting — see the design doc for why.
Per-user "read" state is tracked separately from the notification itself, so
one notification row serves every user rather than being fanned out on create.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
class BaseNotificationStore(ABC):
"""Port interface for perception notification persistence."""
@abstractmethod
def create(
self,
*,
event_id: str,
kind: str,
title: str,
impact_level: str | None,
summary: str | None,
) -> None:
"""Record a new notification. kind is 'new' or 'changed'."""
@abstractmethod
def list_for_user(self, user_id: str, limit: int = 20) -> list[dict]:
"""Return the most recent notifications, newest first.
Each item includes a "read" boolean reflecting whether `user_id` has
marked it read.
"""
@abstractmethod
def unread_count(self, user_id: str) -> int:
"""Return how many notifications `user_id` has not yet read."""
@abstractmethod
def mark_all_read(self, user_id: str) -> int:
"""Mark every currently-unread notification read for `user_id`.
Returns the number of notifications newly marked.
"""
@@ -0,0 +1,43 @@
"""Shared utility functions for crawlers."""
from __future__ import annotations
import re
from datetime import date
def parse_date(text: str) -> str:
"""Return YYYY-MM-DD from common Chinese date formats, or today's date."""
text = text.strip()
if not text:
return date.today().isoformat()
m = re.search(r"(\d{4})[/-](\d{1,2})[/-](\d{1,2})", text)
if m:
try:
return date(int(m.group(1)), int(m.group(2)), int(m.group(3))).isoformat()
except ValueError:
pass
m2 = re.search(r"(\d{4})年(\d{1,2})月(\d{1,2})日?", text)
if m2:
try:
return date(int(m2.group(1)), int(m2.group(2)), int(m2.group(3))).isoformat()
except ValueError:
pass
return date.today().isoformat()
def extract_tags(standard_code: str, title: str) -> list[str]:
"""Derive simple keyword tags from standard code and title."""
tags: list[str] = []
code_upper = standard_code.upper()
if "GB" in code_upper:
tags.append("国家标准")
if "/T" in code_upper:
tags.append("推荐性")
else:
tags.append("强制性")
keywords = ["电动", "安全", "自动驾驶", "充电", "智能网联", "碰撞", "排放", "网络安全"]
for kw in keywords:
if kw in title:
tags.append(kw)
return tags[:5]
@@ -0,0 +1,75 @@
"""Shared contracts for regulatory source crawlers."""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
import httpx
import trafilatura
from loguru import logger
from app.config.settings import settings
@dataclass
class RawEvent:
"""Raw regulatory event returned by a crawler before enrichment."""
source: str
source_label: str
standard_code: str
title: str
summary: str
full_text_url: str
status: str # 'enacted' | 'draft' | 'consultation'
published_at: str # YYYY-MM-DD string
effective_at: str | None
category: str
tags: list[str] = field(default_factory=list)
# Whatever text the list page yields. CrawlService upgrades this by calling
# fetch_full_text(full_text_url); this value is the fallback when that
# fails. Used for change hashing and for the version diff.
raw_text: str = ""
class BaseCrawler(ABC):
"""Abstract regulatory source crawler."""
@abstractmethod
def fetch(self, limit: int = 50) -> list[RawEvent]:
"""Fetch up to `limit` recent events from the data source."""
def fetch_full_text(self, url: str) -> str:
"""Download a regulation detail page and extract its body text.
Change detection is only as good as the text it compares, and list
pages carry nothing but a standard code and a title. This default
implementation serves all current sources; a source that needs PDF
extraction or authentication overrides this one method.
Returns an empty string on any failure rather than raising, so one
unreachable page cannot abort a whole crawl run. The caller decides how
to degrade.
"""
if not url:
return ""
try:
response = httpx.get(
url,
timeout=settings.perception_crawl_timeout_seconds,
follow_redirects=True,
)
response.raise_for_status()
except Exception as exc: # noqa: BLE001 - any transport error degrades the same way
logger.warning("Full-text fetch failed url={} err={}", url, exc)
return ""
# trafilatura scores 0.92 F1 on government pages against 0.78 for
# readability-lxml, and handles CJK content; include_tables matters
# because regulatory limits are frequently tabulated.
extracted = trafilatura.extract(response.text, include_tables=True)
if not extracted:
logger.warning("Full-text extraction returned nothing url={}", url)
return ""
return extracted.strip()
@@ -0,0 +1,88 @@
"""Crawler for CATARC automotive standard catalogue."""
from __future__ import annotations
from urllib.parse import urljoin
import httpx
from bs4 import BeautifulSoup
from loguru import logger
from app.config.settings import settings
from app.infrastructure.perception.crawlers.base import BaseCrawler, RawEvent
from ._utils import extract_tags, parse_date
_BASE_URL = "https://www.catarc.org.cn/bzzxd/qcbz/index.html"
_HOST = "https://www.catarc.org.cn"
_STATUS_MAP = {
"现行": "enacted",
"即将实施": "enacted",
"废止": "enacted",
"征求意见": "consultation",
"报批": "draft",
}
class CatarcCrawler(BaseCrawler):
"""Scrape the CATARC automotive standard list page."""
def fetch(self, limit: int = 50) -> list[RawEvent]:
events: list[RawEvent] = []
page = 1
max_pages = max(10, limit)
while len(events) < limit and page <= max_pages:
url = f"{_BASE_URL}?page={page}"
try:
resp = httpx.get(
url,
timeout=settings.perception_crawl_timeout_seconds,
follow_redirects=True,
)
resp.raise_for_status()
except Exception as exc:
logger.warning("CATARC fetch failed page={} err={}", page, exc)
break
soup = BeautifulSoup(resp.text, "lxml")
rows = soup.select("table tr")
if not rows:
break
batch: list[RawEvent] = []
for row in rows:
cells = row.find_all("td")
if len(cells) < 3:
continue
link = cells[0].find("a")
standard_code = link.get_text(strip=True) if link else cells[0].get_text(strip=True)
title = cells[1].get_text(strip=True) if len(cells) > 1 else standard_code
date_text = cells[2].get_text(strip=True) if len(cells) > 2 else ""
published_at = parse_date(date_text)
status_text = cells[3].get_text(strip=True) if len(cells) > 3 else ""
status = _STATUS_MAP.get(status_text, "enacted")
detail_url = urljoin(_HOST, link["href"]) if link and link.get("href") else url
raw_text = f"{standard_code} {title}"
batch.append(RawEvent(
source="CATARC",
source_label="全国汽车标准化技术委员会",
standard_code=standard_code,
title=title,
summary=title,
full_text_url=detail_url,
status=status,
published_at=published_at,
effective_at=None,
category="汽车标准",
tags=extract_tags(standard_code, title),
raw_text=raw_text,
))
if not batch:
break
events.extend(batch)
page += 1
return events[:limit]
@@ -0,0 +1,122 @@
"""Crawler for EUR-Lex RSS feeds covering EU AI Act and automotive regulations."""
from __future__ import annotations
import re
from email.utils import parsedate_to_datetime
import httpx
from bs4 import BeautifulSoup
from loguru import logger
from app.config.settings import settings
from app.infrastructure.perception.crawlers.base import BaseCrawler, RawEvent
from ._utils import parse_date
_EURLEX_RSS_URLS = [
"https://eur-lex.europa.eu/rss-feed/OJ-L.rss",
]
_AUTOMOTIVE_KEYWORDS = [
"vehicle", "automotive", "motor", "tyre", "emission", "ADAS", "autonomous",
"AI Act", "artificial intelligence", "cybersecurity", "software update",
"R155", "R156", "汽车", "车辆",
]
_AUTOMOTIVE_KEYWORDS_LOWER = [kw.lower() for kw in _AUTOMOTIVE_KEYWORDS]
def _is_automotive_relevant(title: str, description: str) -> bool:
combined = (title + " " + description).lower()
return any(kw in combined for kw in _AUTOMOTIVE_KEYWORDS_LOWER)
def _extract_celex(url: str) -> str:
m = re.search(r"CELEX[:/]([0-9A-Z]+)", url)
return m.group(1) if m else ""
def _parse_rss_date(rfc2822: str) -> str:
try:
dt = parsedate_to_datetime(rfc2822)
return dt.date().isoformat()
except Exception:
return parse_date(rfc2822)
class EurlexCrawler(BaseCrawler):
"""Fetch automotive-relevant EU regulations from EUR-Lex RSS feeds."""
def fetch(self, limit: int = 50) -> list[RawEvent]:
events: list[RawEvent] = []
for rss_url in _EURLEX_RSS_URLS:
if len(events) >= limit:
break
try:
resp = httpx.get(
rss_url,
timeout=settings.perception_crawl_timeout_seconds,
follow_redirects=True,
)
resp.raise_for_status()
except Exception as exc:
logger.warning("EUR-Lex RSS fetch failed url={} err={}", rss_url, exc)
continue
soup = BeautifulSoup(resp.content, "lxml-xml")
for item in soup.find_all("item"):
if len(events) >= limit:
break
title_tag = item.find("title")
title = title_tag.get_text(strip=True) if title_tag else ""
desc_tag = item.find("description")
description = desc_tag.get_text(strip=True) if desc_tag else ""
link_tag = item.find("link")
link = link_tag.get_text(strip=True) if link_tag else ""
pub_date_tag = item.find("pubDate")
pub_date = pub_date_tag.get_text(strip=True) if pub_date_tag else ""
if not _is_automotive_relevant(title, description):
continue
celex = _extract_celex(link)
standard_code = celex if celex else title[:60]
published_at = _parse_rss_date(pub_date) if pub_date else ""
events.append(RawEvent(
source="EUR-Lex",
source_label="欧盟官方公报",
standard_code=standard_code,
title=title,
summary=description[:500],
full_text_url=link,
status="enacted",
published_at=published_at,
effective_at=None,
category="EU法规",
tags=_extract_eurlex_tags(title, description),
raw_text=f"{title}\n{description}",
))
return events[:limit]
def _extract_eurlex_tags(title: str, description: str) -> list[str]:
combined = title + " " + description
tag_map = {
"AI Act": "EU AI Act",
"artificial intelligence": "EU AI Act",
"R155": "UN R155",
"R156": "UN R156",
"cybersecurity": "网络安全",
"emission": "排放",
"autonomous": "自动驾驶",
"ADAS": "ADAS",
}
combined_lower = combined.lower()
tags = []
for kw, tag in tag_map.items():
if kw.lower() in combined_lower:
tags.append(tag)
return tags[:5]
@@ -0,0 +1,98 @@
"""Crawlers for the 国标委 (SAMR) standard information platform."""
from __future__ import annotations
import httpx
from loguru import logger
from app.config.settings import settings
from app.infrastructure.perception.crawlers.base import BaseCrawler, RawEvent
from ._utils import extract_tags, parse_date
_BASE_URL = "https://openstd.samr.gov.cn/bzgk/std/std_list_type"
_HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; RegulatoryBot/1.0)"}
def _fetch_page(std_type: int, page: int, page_size: int) -> list[dict]:
params = {
"p.p1": std_type,
"p.p2": "",
"p.p90": "circulation_date",
"p.p91": "desc",
"p.p6": page,
"p.p7": page_size,
}
try:
resp = httpx.get(
_BASE_URL,
params=params,
headers=_HEADERS,
timeout=settings.perception_crawl_timeout_seconds,
)
resp.raise_for_status()
data = resp.json()
return data.get("rows", []) or []
except Exception as exc:
logger.warning("国标委 fetch failed type={} page={} err={}", std_type, page, exc)
return []
def _row_to_raw_event(row: dict, source_label: str) -> RawEvent:
standard_code = row.get("std_code", "")
title = row.get("std_name", standard_code)
published_at = parse_date(row.get("release_date", ""))
effective_at_raw = row.get("implement_date", "")
effective_at = parse_date(effective_at_raw) if effective_at_raw else None
status_text = row.get("std_status", "")
if "征求意见" in status_text:
status = "consultation"
elif "报批" in status_text or "草案" in status_text:
status = "draft"
else:
status = "enacted"
return RawEvent(
source="国标委",
source_label=source_label,
standard_code=standard_code,
title=title,
summary=title,
full_text_url=f"https://openstd.samr.gov.cn/bzgk/std/detail?id={row.get('id', '')}",
status=status,
published_at=published_at,
effective_at=effective_at,
category=row.get("std_type", "国家标准"),
tags=extract_tags(standard_code, title),
raw_text=f"{standard_code} {title}",
)
class GuobiaoMandatoryCrawler(BaseCrawler):
"""Fetch mandatory national standards (强制性) related to vehicles."""
def fetch(self, limit: int = 50) -> list[RawEvent]:
events: list[RawEvent] = []
page = 1
max_pages = max(10, limit)
while len(events) < limit and page <= max_pages:
rows = _fetch_page(std_type=1, page=page, page_size=20)
if not rows:
break
events.extend(_row_to_raw_event(r, "国标委·强制性") for r in rows)
page += 1
return events[:limit]
class GuobiaoRecommendedCrawler(BaseCrawler):
"""Fetch recommended national standards (推荐性) related to vehicles."""
def fetch(self, limit: int = 50) -> list[RawEvent]:
events: list[RawEvent] = []
page = 1
max_pages = max(10, limit)
while len(events) < limit and page <= max_pages:
rows = _fetch_page(std_type=2, page=page, page_size=20)
if not rows:
break
events.extend(_row_to_raw_event(r, "国标委·推荐性") for r in rows)
page += 1
return events[:limit]
@@ -0,0 +1,245 @@
"""LLM-driven pipeline for regulatory event enrichment."""
from __future__ import annotations
import json
from typing import Any
from loguru import logger
from app.config.settings import settings
from app.infrastructure.perception.regulation_differ import (
ParagraphChange,
RegulationDiffer,
)
from app.services.llm.llm_factory import get_llm_client
_EXTRACT_SYSTEM = (
"You are a regulatory compliance expert specialising in automotive standards "
"(GB, UN-ECE, ISO, EU). Extract structured information from regulation text. "
"Return valid JSON only — no markdown fences, no extra keys."
)
_ASSESS_SYSTEM = (
"You are an automotive compliance analyst. Given a regulation and related document excerpts, "
"identify which documents are affected and what actions are required. "
"Return a JSON array only."
)
_DIFF_SYSTEM = (
"You are a regulatory change analyst. You are given the OLD and NEW version of "
"one regulation paragraph, with the exact edits marked <DEL>removed</DEL> and "
"<INS>added</INS>. Classify the legal effect of the change. "
"Return JSON only: {\"change_type\": \"tightened|relaxed|numeric|clarified|scope\", "
"\"legal_effect\": \"one sentence on what this means for compliance\"}"
)
def _marked_diff(change: ParagraphChange) -> str:
"""Render a paragraph change with the exact edits marked for the model.
The model is shown where the edit is rather than being asked to find it,
and is never asked to reproduce the changed text — the differ already
computed those spans exactly, so there is nothing for the model to
hallucinate.
"""
marked = "".join(
text if op == 0 else (f"<DEL>{text}</DEL>" if op < 0 else f"<INS>{text}</INS>")
for op, text in change.diff_spans
)
return (
f"OLD: {change.old_text[:500]}\n"
f"NEW: {change.new_text[:500]}\n"
f"MARKED: {marked[:800]}"
)
def _llm_json(client: Any, messages: list[dict]) -> Any:
"""Call LLM and parse JSON response; return None on failure."""
try:
resp = client.chat(messages)
text = (resp.content or "").strip()
if text.startswith("```"):
text = text.split("```")[1]
if text.startswith("json"):
text = text[4:]
return json.loads(text)
except Exception as exc:
logger.warning("LLM JSON parse failed: {}", exc)
return None
class LlmPipeline:
"""Three-step enrichment pipeline for crawled regulatory events."""
def __init__(self) -> None:
self._client = get_llm_client(
provider=settings.llm_provider,
model=settings.llm_model,
)
# Change detection is deterministic; the differ needs no model and no
# network, so the pipeline no longer constructs an embedding provider.
self._differ = RegulationDiffer()
# ------------------------------------------------------------------
# Step 1: Structure extraction
# ------------------------------------------------------------------
def extract_structure(self, event: dict) -> dict:
"""Extract obligations, deadlines, scope, penalties, impact_level from event text."""
prompt = f"""Extract structured compliance information from this regulation:
Standard: {event.get('standard_code', '')}
Title: {event.get('title', '')}
Source: {event.get('source_label', '')}
Summary: {event.get('summary', '')}
Tags: {', '.join(event.get('tags') or [])}
Return JSON with exactly these keys:
{{
"obligations": [{{"text": "...", "deontic": "must|shall|may|prohibited", "subject": "...", "object": "...", "condition": ""}}],
"deadlines": [{{"date": "YYYY-MM-DD or null", "description": "..."}}],
"scope": "one sentence describing who/what this applies to",
"penalties": "one sentence on consequences of non-compliance, or null",
"impact_level": "high|medium|low"
}}"""
messages = [
{"role": "system", "content": _EXTRACT_SYSTEM},
{"role": "user", "content": prompt},
]
result = _llm_json(self._client, messages)
if not isinstance(result, dict):
return {
"obligations": [],
"deadlines": [],
"scope": "",
"penalties": "",
"impact_level": "medium",
}
return result
# ------------------------------------------------------------------
# Step 2: Impact assessment
# ------------------------------------------------------------------
def assess_impact(self, event: dict, retrieval_service: Any) -> list[dict]:
"""Use RAG to find affected documents and generate recommendations."""
obligations = event.get("obligations") or []
obligation_texts = " ".join(o.get("text", "") for o in obligations[:3])
query = f"{event.get('standard_code', '')} {event.get('title', '')} {obligation_texts}"
try:
chunks = retrieval_service.retrieve(query=query, top_k=5)
except Exception as exc:
logger.warning("RAG retrieval failed: {}", exc)
return []
if not chunks:
return []
seen: set[str] = set()
doc_excerpts: list[dict] = []
for chunk in chunks:
if chunk.doc_id not in seen:
seen.add(chunk.doc_id)
doc_excerpts.append({
"doc_id": chunk.doc_id,
"doc_name": chunk.doc_title,
"score": round(float(chunk.score if chunk.score is not None else 0), 4),
"snippet": (chunk.text or "")[:300],
"clause": getattr(chunk, "section_title", "") or "",
})
context = "\n".join(
f"[{d['doc_name']} {d['clause']}] score={d['score']}: {d['snippet']}"
for d in doc_excerpts
)
prompt = f"""Regulation: {event.get('standard_code')}{event.get('title')}
Obligations: {obligation_texts or event.get('summary', '')}
Affected documents found in knowledge base:
{context}
For each document, assess impact and recommend action. Return JSON array:
[{{"doc_id":"...","doc_name":"...","score":0.0,"key_clauses":"...","recommendation":"one sentence action"}}]"""
messages = [
{"role": "system", "content": _ASSESS_SYSTEM},
{"role": "user", "content": prompt},
]
result = _llm_json(self._client, messages)
if isinstance(result, list):
score_map = {d["doc_id"]: d["score"] for d in doc_excerpts}
for item in result:
if isinstance(item, dict) and item.get("doc_id") in score_map:
item["score"] = score_map[item["doc_id"]]
return result
return doc_excerpts
# ------------------------------------------------------------------
# Step 3: Deterministic diff with gated LLM classification
# ------------------------------------------------------------------
def compute_diff(self, old_text: str, new_text: str) -> dict:
"""Compare old and new regulation text; return changed sections and summary.
Detection is deterministic — see regulation_differ for why embedding
similarity was removed. The LLM is called only for paragraphs the
differ marked significant, and only to explain the legal effect of a
change that has already been located exactly.
"""
changes = self._differ.diff(old_text, new_text)
if not changes:
return {
"changed_sections": [],
"change_summary": "No substantive changes detected between versions.",
}
changed_sections = [self._describe(change) for change in changes]
types = sorted({section["change_type"] for section in changed_sections})
gated = sum(1 for change in changes if change.needs_llm)
change_summary = (
f"{len(changed_sections)} paragraph(s) changed ({', '.join(types)}); "
f"{gated} significant. "
+ (changed_sections[0].get("summary") or "")
).strip()
return {"changed_sections": changed_sections, "change_summary": change_summary}
def _describe(self, change: ParagraphChange) -> dict:
"""Turn one detected change into the API payload, classifying if warranted."""
section = {
"old_text": change.old_text[:300],
"new_text": change.new_text[:300],
"change_type": change.change_type,
"change_ratio": round(change.change_ratio, 3),
"numeric_changed": change.numeric_changed,
"deontic_changed": change.deontic_changed,
"summary": "",
}
if not change.needs_llm:
return section
classification = _llm_json(
self._client,
[
{"role": "system", "content": _DIFF_SYSTEM},
{"role": "user", "content": _marked_diff(change)},
],
)
if isinstance(classification, dict):
section["change_type"] = classification.get("change_type") or change.change_type
section["summary"] = classification.get("legal_effect") or ""
# A failed or malformed model response must not discard a change that
# deterministic analysis already proved real; the section keeps its
# spans, flags, and alignment-derived type with an empty summary.
if change.numeric_changed:
# Models routinely label a changed threshold as "clarified". The
# deterministic pass already knows a number moved, so it wins.
section["change_type"] = "numeric"
return section
@@ -4,6 +4,8 @@ from __future__ import annotations
from typing import Any
from app.infrastructure.perception.base_event_store import BaseEventStore
MOCK_EVENTS: list[dict[str, Any]] = [
# ------------------------------------------------------------------ HIGH
{
@@ -379,18 +381,18 @@ MOCK_EVENTS: list[dict[str, Any]] = [
},
]
# Index for fast lookup
_EVENT_INDEX: dict[str, dict] = {e["id"]: e for e in MOCK_EVENTS}
class MockEventStore:
class MockEventStore(BaseEventStore):
"""In-memory mock store for regulatory events."""
def __init__(self) -> None:
self._events: list[dict] = [dict(e) for e in MOCK_EVENTS]
self._index: dict[str, dict] = {e["id"]: e for e in self._events}
def all(self) -> list[dict]:
return list(MOCK_EVENTS)
return list(self._events)
def get(self, event_id: str) -> dict | None:
return _EVENT_INDEX.get(event_id)
return self._index.get(event_id)
def filter(
self,
@@ -399,23 +401,39 @@ class MockEventStore:
impact_level: str | None = None,
limit: int = 50,
) -> list[dict]:
events = list(MOCK_EVENTS)
events = list(self._events)
if source:
events = [e for e in events if e["source"] == source]
if impact_level:
events = [e for e in events if e["impact_level"] == impact_level]
events.sort(key=lambda e: e["published_at"], reverse=True)
events.sort(key=lambda e: e.get("published_at") or "", reverse=True)
return events[:limit]
def stats(self) -> dict:
from datetime import date, timedelta
events = MOCK_EVENTS
events = self._events
cutoff = (date.today() - timedelta(days=90)).isoformat()
return {
"total": len(events),
"high_impact": sum(1 for e in events if e["impact_level"] == "high"),
"medium_impact": sum(1 for e in events if e["impact_level"] == "medium"),
"low_impact": sum(1 for e in events if e["impact_level"] == "low"),
"recent_90d": sum(1 for e in events if e["published_at"] >= cutoff),
"recent_90d": sum(1 for e in events if (e.get("published_at") or "") >= cutoff),
}
def upsert(self, event: dict) -> None:
"""Insert or update event in the in-memory list (used in tests)."""
existing = self._index.get(event["id"])
if existing:
existing.update(event)
else:
self._events.append(event)
self._index[event["id"]] = event
def get_by_standard_code(self, standard_code: str) -> dict | None:
"""Return most-recent event with matching standard_code."""
matches = [e for e in self._events if e.get("standard_code") == standard_code]
if not matches:
return None
return max(matches, key=lambda e: e.get("published_at", ""))
@@ -0,0 +1,66 @@
"""In-memory notification store used when Postgres is not configured.
Mirrors MockEventStore's role for BaseEventStore: keeps the feature usable in
local dev and in tests without a live database, and matches
DOCUMENT_REPOSITORY_BACKEND's existing Mock/Postgres split.
"""
from __future__ import annotations
from datetime import UTC, datetime
from app.infrastructure.perception.base_notification_store import BaseNotificationStore
class MockNotificationStore(BaseNotificationStore):
"""Dict-backed notification store. Data does not survive a process restart."""
def __init__(self) -> None:
"""Start with an empty feed and no read receipts."""
self._notifications: list[dict] = []
self._next_id = 1
# (notification_id, user_id) pairs — presence means read.
self._reads: set[tuple[int, str]] = set()
def create(
self,
*,
event_id: str,
kind: str,
title: str,
impact_level: str | None,
summary: str | None,
) -> None:
"""Append a notification with an auto-incrementing id."""
self._notifications.append({
"id": self._next_id,
"event_id": event_id,
"kind": kind,
"title": title,
"impact_level": impact_level,
"summary": summary,
"created_at": datetime.now(UTC).isoformat(),
})
self._next_id += 1
def list_for_user(self, user_id: str, limit: int = 20) -> list[dict]:
"""Return the newest `limit` notifications with this user's read state."""
ordered = sorted(self._notifications, key=lambda n: n["id"], reverse=True)
return [
{**n, "read": (n["id"], user_id) in self._reads}
for n in ordered[:limit]
]
def unread_count(self, user_id: str) -> int:
"""Count notifications this user has not yet read."""
return sum(1 for n in self._notifications if (n["id"], user_id) not in self._reads)
def mark_all_read(self, user_id: str) -> int:
"""Add a read receipt for every currently-unread notification."""
marked = 0
for n in self._notifications:
key = (n["id"], user_id)
if key not in self._reads:
self._reads.add(key)
marked += 1
return marked
@@ -0,0 +1,234 @@
"""PostgreSQL-backed regulatory event store."""
from __future__ import annotations
import json
from contextlib import contextmanager
from datetime import UTC, date, datetime, timedelta
from typing import Any
import psycopg2
import psycopg2.extras
from psycopg2.pool import ThreadedConnectionPool
from app.config.settings import settings
from app.infrastructure.perception.base_event_store import BaseEventStore
_CREATE_TABLE = """
CREATE TABLE IF NOT EXISTS regulation_events (
id TEXT PRIMARY KEY,
source TEXT NOT NULL,
source_label TEXT,
standard_code TEXT NOT NULL,
title TEXT NOT NULL,
summary TEXT,
full_text_url TEXT,
status TEXT,
impact_level TEXT,
published_at DATE,
effective_at DATE,
category TEXT,
tags TEXT[],
obligations JSONB,
deadlines JSONB,
scope TEXT,
penalties TEXT,
content_hash TEXT,
previous_hash TEXT,
change_summary TEXT,
changed_sections JSONB,
affected_docs JSONB,
crawled_at TIMESTAMPTZ DEFAULT now(),
processed_at TIMESTAMPTZ,
raw_storage_key TEXT,
raw_text TEXT
);
CREATE INDEX IF NOT EXISTS reg_events_source_date
ON regulation_events (source, published_at DESC);
CREATE INDEX IF NOT EXISTS reg_events_impact_date
ON regulation_events (impact_level, published_at DESC);
"""
_ADD_COLUMNS = """
ALTER TABLE regulation_events ADD COLUMN IF NOT EXISTS raw_text TEXT;
"""
_ALL_COLUMNS = (
"id", "source", "source_label", "standard_code", "title", "summary",
"full_text_url", "status", "impact_level", "published_at", "effective_at",
"category", "tags", "obligations", "deadlines", "scope", "penalties",
"content_hash", "previous_hash", "change_summary", "changed_sections",
"affected_docs", "crawled_at", "processed_at", "raw_storage_key", "raw_text",
)
def _row_to_dict(row: dict[str, Any]) -> dict:
"""Convert a psycopg2 RealDictRow to a plain dict with serialized JSON fields."""
d = dict(row)
for field in ("obligations", "deadlines", "changed_sections", "affected_docs"):
val = d.get(field)
if isinstance(val, str):
d[field] = json.loads(val)
for date_field in ("published_at", "effective_at"):
val = d.get(date_field)
if isinstance(val, datetime):
d[date_field] = val.date().isoformat()
elif isinstance(val, date):
d[date_field] = val.isoformat()
for ts_field in ("crawled_at", "processed_at"):
val = d.get(ts_field)
if isinstance(val, datetime):
d[ts_field] = val.isoformat()
return d
class PostgresEventStore(BaseEventStore):
"""Regulatory event store backed by PostgreSQL."""
def __init__(self) -> None:
self._pool = ThreadedConnectionPool(
minconn=1,
maxconn=5,
host=settings.postgres_host,
port=settings.postgres_port,
user=settings.postgres_user,
password=settings.postgres_password,
dbname=settings.postgres_db,
)
self._ensure_schema()
def _ensure_schema(self) -> None:
with self._conn() as conn:
try:
with conn.cursor() as cur:
cur.execute(_CREATE_TABLE)
# CREATE TABLE IF NOT EXISTS is a no-op on deployments that
# already have this table, so new columns must be added
# explicitly or existing installations silently lack them.
cur.execute(_ADD_COLUMNS)
conn.commit()
except Exception:
conn.rollback()
raise
@contextmanager
def _conn(self):
conn = None
try:
conn = self._pool.getconn()
yield conn
finally:
if conn is not None:
self._pool.putconn(conn)
def all(self) -> list[dict]:
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"SELECT * FROM regulation_events ORDER BY published_at DESC NULLS LAST"
)
return [_row_to_dict(r) for r in cur.fetchall()]
def get(self, event_id: str) -> dict | None:
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"SELECT * FROM regulation_events WHERE id = %s", (event_id,)
)
row = cur.fetchone()
return _row_to_dict(row) if row else None
def filter(
self,
*,
source: str | None = None,
impact_level: str | None = None,
limit: int = 50,
) -> list[dict]:
conditions: list[str] = []
params: list[Any] = []
if source:
conditions.append("source = %s")
params.append(source)
if impact_level:
conditions.append("impact_level = %s")
params.append(impact_level)
where = ("WHERE " + " AND ".join(conditions)) if conditions else ""
params.append(limit)
sql = f"""
SELECT * FROM regulation_events
{where}
ORDER BY published_at DESC NULLS LAST
LIMIT %s
"""
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(sql, params)
return [_row_to_dict(r) for r in cur.fetchall()]
def stats(self) -> dict:
cutoff = (date.today() - timedelta(days=90)).isoformat()
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute("SELECT COUNT(*) AS count FROM regulation_events")
total = (cur.fetchone() or {}).get("count", 0)
cur.execute(
"SELECT COUNT(*) AS count FROM regulation_events WHERE impact_level = 'high'"
)
high = (cur.fetchone() or {}).get("count", 0)
cur.execute(
"SELECT COUNT(*) AS count FROM regulation_events WHERE impact_level = 'medium'"
)
medium = (cur.fetchone() or {}).get("count", 0)
cur.execute(
"SELECT COUNT(*) AS count FROM regulation_events WHERE published_at >= %s",
(cutoff,),
)
recent = (cur.fetchone() or {}).get("count", 0)
return {
"total": int(total),
"high_impact": int(high),
"medium_impact": int(medium),
"recent_90d": int(recent),
}
def upsert(self, event: dict) -> None:
"""Insert or update a regulation event."""
cols = [c for c in _ALL_COLUMNS if c in event]
placeholders = ", ".join(f"%({c})s" for c in cols)
updates = ", ".join(f"{c} = EXCLUDED.{c}" for c in cols if c != "id")
sql = f"""
INSERT INTO regulation_events ({', '.join(cols)})
VALUES ({placeholders})
ON CONFLICT (id) DO UPDATE SET {updates}
"""
row: dict[str, Any] = {}
for c in cols:
val = event.get(c)
if c in ("obligations", "deadlines", "changed_sections", "affected_docs") and val is not None:
row[c] = json.dumps(val, ensure_ascii=False)
elif c == "tags" and isinstance(val, list):
row[c] = val
else:
row[c] = val
with self._conn() as conn:
try:
with conn.cursor() as cur:
cur.execute(sql, row)
conn.commit()
except Exception:
conn.rollback()
raise
def get_by_standard_code(self, standard_code: str) -> dict | None:
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""SELECT * FROM regulation_events
WHERE standard_code = %s
ORDER BY published_at DESC NULLS LAST
LIMIT 1""",
(standard_code,),
)
row = cur.fetchone()
return _row_to_dict(row) if row else None
@@ -0,0 +1,157 @@
"""PostgreSQL-backed notification store.
One row per triggering event, shared by every user; a separate read-receipt
table tracks per-user read state so broadcasting to everyone needs no fan-out
insert per user. See base_notification_store.py for the port contract and the
design doc for why this shape was chosen over per-user subscriptions.
"""
from __future__ import annotations
from contextlib import contextmanager
from typing import Any
import psycopg2
import psycopg2.extras
from psycopg2.pool import ThreadedConnectionPool
from app.config.settings import settings
from app.infrastructure.perception.base_notification_store import BaseNotificationStore
_CREATE_TABLES = """
CREATE TABLE IF NOT EXISTS perception_notifications (
id SERIAL PRIMARY KEY,
event_id TEXT NOT NULL REFERENCES regulation_events(id) ON DELETE CASCADE,
kind TEXT NOT NULL,
title TEXT NOT NULL,
impact_level TEXT,
summary TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS perception_notification_reads (
notification_id INTEGER NOT NULL REFERENCES perception_notifications(id) ON DELETE CASCADE,
user_id TEXT NOT NULL,
read_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (notification_id, user_id)
);
CREATE INDEX IF NOT EXISTS perception_notif_created
ON perception_notifications (created_at DESC);
"""
def _row_to_dict(row: dict[str, Any]) -> dict:
"""Convert a psycopg2 RealDictRow to a plain dict with an ISO timestamp."""
d = dict(row)
if d.get("created_at") is not None:
d["created_at"] = d["created_at"].isoformat()
return d
class PostgresNotificationStore(BaseNotificationStore):
"""Notification store backed by PostgreSQL."""
def __init__(self) -> None:
"""Open a connection pool and ensure both tables exist."""
self._pool = ThreadedConnectionPool(
minconn=1,
maxconn=5,
host=settings.postgres_host,
port=settings.postgres_port,
user=settings.postgres_user,
password=settings.postgres_password,
dbname=settings.postgres_db,
)
self._ensure_schema()
def _ensure_schema(self) -> None:
with self._conn() as conn:
try:
with conn.cursor() as cur:
cur.execute(_CREATE_TABLES)
conn.commit()
except Exception:
conn.rollback()
raise
@contextmanager
def _conn(self):
conn = None
try:
conn = self._pool.getconn()
yield conn
finally:
if conn is not None:
self._pool.putconn(conn)
def create(
self,
*,
event_id: str,
kind: str,
title: str,
impact_level: str | None,
summary: str | None,
) -> None:
"""Insert one notification row for the triggering event."""
with self._conn() as conn:
with conn.cursor() as cur:
cur.execute(
"INSERT INTO perception_notifications "
"(event_id, kind, title, impact_level, summary) "
"VALUES (%s, %s, %s, %s, %s)",
(event_id, kind, title, impact_level, summary),
)
conn.commit()
def list_for_user(self, user_id: str, limit: int = 20) -> list[dict]:
"""Return the newest notifications with this user's read state joined in."""
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
SELECT n.*, (r.user_id IS NOT NULL) AS read
FROM perception_notifications n
LEFT JOIN perception_notification_reads r
ON r.notification_id = n.id AND r.user_id = %s
ORDER BY n.created_at DESC
LIMIT %s
""",
(user_id, limit),
)
return [_row_to_dict(r) for r in cur.fetchall()]
def unread_count(self, user_id: str) -> int:
"""Count notifications with no read receipt for this user."""
with self._conn() as conn:
with conn.cursor() as cur:
cur.execute(
"""
SELECT COUNT(*) FROM perception_notifications n
WHERE NOT EXISTS (
SELECT 1 FROM perception_notification_reads r
WHERE r.notification_id = n.id AND r.user_id = %s
)
""",
(user_id,),
)
return cur.fetchone()[0]
def mark_all_read(self, user_id: str) -> int:
"""Insert a read receipt for every notification this user hasn't read."""
with self._conn() as conn:
with conn.cursor() as cur:
cur.execute(
"""
INSERT INTO perception_notification_reads (notification_id, user_id)
SELECT n.id, %s FROM perception_notifications n
WHERE NOT EXISTS (
SELECT 1 FROM perception_notification_reads r
WHERE r.notification_id = n.id AND r.user_id = %s
)
ON CONFLICT (notification_id, user_id) DO NOTHING
""",
(user_id, user_id),
)
marked = cur.rowcount
conn.commit()
return marked
@@ -0,0 +1,224 @@
"""Deterministic change detection between two versions of a regulation.
This module deliberately contains no LLM call, no network access, and no
embedding lookup. It exists because the previous implementation decided whether
a paragraph had changed by comparing embedding cosine similarity against a 0.85
threshold, which is blind to exactly the edits that matter in regulation.
Measured against the deployed text-embedding-v3 gateway, tightening a braking
limit from 30米 to 20米 scores 0.9153 and relaxing 应当 to 宜 scores 0.9162 —
both far above the threshold, both undetected — while an entirely unrelated
clause scores 0.6862 and is the only thing that fires. Cosine is scale
invariant, so it cannot represent a change in magnitude or certainty
(arXiv:2403.05440, ACM Web Conference 2024); no threshold recovers the signal.
The replacement is the production consensus for legal text: align paragraphs
with a longest-common-subsequence matcher, run a literal character diff on the
aligned pairs, and let cheap deterministic rules decide whether a change is
significant enough to spend an LLM call classifying.
"""
from __future__ import annotations
import re
import unicodedata
from dataclasses import dataclass, field
from diff_match_patch import diff_match_patch
from difflib import SequenceMatcher
from app.config.settings import settings
# Chinese regulatory drafting uses a small, near-unambiguous set of deontic
# markers, so a regex pre-pass identifies legally significant edits without an
# LLM. Adding or removing any of these changes what the provision compels.
_DEONTIC_PATTERN = re.compile(r"应当|须|禁止|不得|可以|允许|宜")
# Matches digit runs including decimals, so "30" -> "20" and "0.85" -> "0.9"
# are both treated as numeric changes.
_NUMBER_PATTERN = re.compile(r"\d+(?:\.\d+)?")
# diff_match_patch operation codes.
_DMP_DELETE = -1
_DMP_INSERT = 1
_DMP_EQUAL = 0
@dataclass(frozen=True)
class ParagraphChange:
"""One detected difference between the old and new version of a regulation.
`needs_llm` is the gate: it records whether this change is worth the cost of
an LLM classification call. The deterministic flags that drive it are kept
on the record so downstream code can act on them even when the LLM call
fails or is skipped.
"""
change_type: str
old_text: str
new_text: str
numeric_changed: bool
deontic_changed: bool
change_ratio: float
needs_llm: bool
# (op, text) pairs from diff_match_patch, for rendering a redline view.
diff_spans: list[tuple[int, str]] = field(default_factory=list)
def _split_paragraphs(text: str) -> list[str]:
"""Split regulation text into comparable units, dropping blank lines.
ponytail: newline splitting, not clause parsing. Upgrade to 第X条 / X.X.X
segmentation only if paragraph granularity proves too coarse in practice.
"""
return [line.strip() for line in (text or "").split("\n") if line.strip()]
def _numbers_differ(old: str, new: str) -> bool:
"""Report whether the two spans contain a different sequence of numbers."""
return _NUMBER_PATTERN.findall(old) != _NUMBER_PATTERN.findall(new)
def _deontic_differs(old: str, new: str) -> bool:
"""Report whether obligation markers were added, removed, or swapped."""
return sorted(_DEONTIC_PATTERN.findall(old)) != sorted(_DEONTIC_PATTERN.findall(new))
def _is_cosmetic(spans: list[tuple[int, str]]) -> bool:
"""Report whether the edit touched nothing but punctuation and whitespace.
A change ratio alone cannot answer this for Chinese regulation text. Clauses
run 20-60 characters, so deleting a single 。 is a 4% change and clears any
threshold low enough to still catch real edits in longer paragraphs. Testing
what actually changed is both cheaper and exact.
"""
changed = "".join(text for op, text in spans if op != _DMP_EQUAL)
# Unicode categories P (punctuation), Z (separator) and C (control) cover
# Chinese and ASCII punctuation plus every flavour of whitespace.
return all(unicodedata.category(char)[0] in {"P", "Z", "C"} for char in changed)
class RegulationDiffer:
"""Align two regulation versions and classify what changed, without an LLM."""
def __init__(self, min_change_ratio: float | None = None) -> None:
"""Store the gate threshold, defaulting to the configured value.
The explicit argument exists so tests never depend on the deployed .env.
"""
self._min_change_ratio = (
settings.perception_diff_min_change_ratio
if min_change_ratio is None
else min_change_ratio
)
self._dmp = diff_match_patch()
def diff(self, old_text: str, new_text: str) -> list[ParagraphChange]:
"""Return every changed paragraph between two versions.
Unchanged paragraphs are not returned. An empty old version means there
is no baseline to compare against — the caller's first crawl — so no
changes are reported rather than the whole document being called new.
"""
old_paras = _split_paragraphs(old_text)
new_paras = _split_paragraphs(new_text)
if not old_paras or not new_paras:
return []
# autojunk=False is load-bearing: the default treats any element
# appearing in over 1% of a sequence of 200+ items as junk, and
# regulations repeat boilerplate paragraphs that alignment depends on
# as anchors.
matcher = SequenceMatcher(None, old_paras, new_paras, autojunk=False)
changes: list[ParagraphChange] = []
for tag, i1, i2, j1, j2 in matcher.get_opcodes():
if tag == "equal":
continue
if tag == "insert":
changes.extend(self._added(p) for p in new_paras[j1:j2])
elif tag == "delete":
changes.extend(self._removed(p) for p in old_paras[i1:i2])
elif tag == "replace":
changes.extend(self._replaced(old_paras[i1:i2], new_paras[j1:j2]))
return changes
def _added(self, paragraph: str) -> ParagraphChange:
"""Build a record for a provision present only in the new version."""
return ParagraphChange(
change_type="added",
old_text="",
new_text=paragraph,
numeric_changed=False,
deontic_changed=bool(_DEONTIC_PATTERN.search(paragraph)),
change_ratio=1.0,
# A new provision always carries new obligations, so it is always
# worth classifying.
needs_llm=True,
diff_spans=[(_DMP_INSERT, paragraph)],
)
def _removed(self, paragraph: str) -> ParagraphChange:
"""Build a record for a provision dropped from the new version."""
return ParagraphChange(
change_type="removed",
old_text=paragraph,
new_text="",
numeric_changed=False,
deontic_changed=bool(_DEONTIC_PATTERN.search(paragraph)),
change_ratio=1.0,
needs_llm=True,
diff_spans=[(_DMP_DELETE, paragraph)],
)
def _replaced(self, old_block: list[str], new_block: list[str]) -> list[ParagraphChange]:
"""Compare a run of rewritten paragraphs pairwise, reporting the remainder.
SequenceMatcher emits `replace` for a whole run at once, and the two
sides may differ in length. Pairing by position within the run is safe
here because alignment has already established that this run as a whole
corresponds; any surplus on either side is a genuine insertion or
deletion.
"""
results: list[ParagraphChange] = []
for index in range(max(len(old_block), len(new_block))):
if index >= len(old_block):
results.append(self._added(new_block[index]))
elif index >= len(new_block):
results.append(self._removed(old_block[index]))
else:
results.append(self._modified(old_block[index], new_block[index]))
return results
def _modified(self, old: str, new: str) -> ParagraphChange:
"""Character-diff an aligned pair and decide whether it warrants an LLM call."""
spans = self._dmp.diff_main(old, new)
# Merges single-character edits into human-meaningful chunks so the
# redline view and the change ratio both reflect real edits.
self._dmp.diff_cleanupSemantic(spans)
changed_chars = sum(len(text) for op, text in spans if op != _DMP_EQUAL)
denominator = max(len(old), len(new), 1)
change_ratio = changed_chars / denominator
numeric_changed = _numbers_differ(old, new)
deontic_changed = _deontic_differs(old, new)
# A changed limit or obligation marker is always significant no matter
# how few characters moved. Everything else must be substantive and
# clear the ratio gate to be worth a model call.
significant = numeric_changed or deontic_changed or (
not _is_cosmetic(spans) and change_ratio >= self._min_change_ratio
)
return ParagraphChange(
change_type="modified",
old_text=old,
new_text=new,
numeric_changed=numeric_changed,
deontic_changed=deontic_changed,
change_ratio=change_ratio,
needs_llm=significant,
diff_spans=[(op, text) for op, text in spans],
)
@@ -0,0 +1,169 @@
"""Redis-backed conversation store for persistent chat sessions.
Sessions are stored as JSON strings under the key `session:{session_id}`.
The Redis TTL is refreshed on every write so active sessions stay alive.
On expiry, `get_session` returns None — callers should create a new session.
"""
from __future__ import annotations
import json
import time
import uuid
from typing import Any
from loguru import logger
from app.domain.conversation import ConversationMessage, ConversationSession, ConversationStore
class RedisConversationStore(ConversationStore):
"""Store conversation sessions in Redis with automatic TTL expiry.
Each session is serialised as a JSON object at key ``session:{session_id}``.
The TTL is reset on every write so sessions stay alive as long as they are active.
"""
# Prefix for all session keys to avoid collisions with other Redis consumers.
_PREFIX = "session:"
def __init__(self, *, redis_client: Any, timeout_seconds: int = 1800) -> None:
"""Initialise the store with an existing Redis client and a TTL in seconds."""
self._redis = redis_client
self._ttl = timeout_seconds
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _key(self, session_id: str) -> str:
"""Build the Redis key for a session."""
return f"{self._PREFIX}{session_id}"
def _serialise(self, session: ConversationSession) -> str:
"""Serialise a ConversationSession to a JSON string."""
return json.dumps(
{
"session_id": session.session_id,
"created_at": session.created_at,
"updated_at": session.updated_at,
"metadata": session.metadata,
"messages": [
{
"role": msg.role,
"content": msg.content,
"timestamp": msg.timestamp,
"sources": msg.sources,
}
for msg in session.messages
],
},
ensure_ascii=False,
)
def _deserialise(self, raw: bytes | str) -> ConversationSession:
"""Deserialise a JSON string back into a ConversationSession."""
data = json.loads(raw)
messages = [
ConversationMessage(
role=m["role"],
content=m["content"],
timestamp=m["timestamp"],
sources=m.get("sources", []),
)
for m in data.get("messages", [])
]
session = ConversationSession(
session_id=data["session_id"],
created_at=data.get("created_at", 0),
updated_at=data.get("updated_at", 0),
metadata=data.get("metadata", {}),
)
session.messages = messages
return session
def _save(self, session: ConversationSession) -> None:
"""Persist a session to Redis and refresh its TTL."""
self._redis.setex(self._key(session.session_id), self._ttl, self._serialise(session))
# ------------------------------------------------------------------
# ConversationStore protocol
# ------------------------------------------------------------------
def create_session(self, metadata: dict | None = None) -> ConversationSession:
"""Create a new empty session and persist it immediately."""
now = int(time.time())
session = ConversationSession(
session_id=str(uuid.uuid4())[:8],
created_at=now,
updated_at=now,
metadata=metadata or {},
)
self._save(session)
return session
def get_session(self, session_id: str) -> ConversationSession | None:
"""Return a session by ID, or None if it does not exist or has expired."""
raw = self._redis.get(self._key(session_id))
if raw is None:
return None
try:
return self._deserialise(raw)
except Exception:
logger.warning("Failed to deserialise session: {}", session_id)
return None
def save_message(
self,
session_id: str,
*,
role: str,
content: str,
sources: list[dict] | None = None,
) -> ConversationSession | None:
"""Append a message to a session and refresh its TTL."""
session = self.get_session(session_id)
if session is None:
return None
session.messages.append(
ConversationMessage(
role=role,
content=content,
timestamp=int(time.time()),
sources=sources or [],
)
)
session.updated_at = int(time.time())
self._save(session)
return session
def delete_session(self, session_id: str) -> bool:
"""Delete a session. Returns True if it existed, False otherwise."""
deleted = self._redis.delete(self._key(session_id))
return bool(deleted)
def list_sessions(self) -> list[dict]:
"""Return summary dicts for all live sessions visible in this Redis DB.
Note: KEYS is used for simplicity; replace with SCAN for large deployments.
"""
pattern = f"{self._PREFIX}*"
keys = self._redis.keys(pattern)
result = []
for key in keys:
raw = self._redis.get(key)
if raw is None:
continue
try:
data = json.loads(raw)
result.append(
{
"session_id": data["session_id"],
"message_count": len(data.get("messages", [])),
"created_at": data.get("created_at", 0),
"updated_at": data.get("updated_at", 0),
}
)
except Exception:
continue
return result
@@ -41,6 +41,10 @@ class MinioDocumentBinaryStore(DocumentBinaryStore):
raise FileNotFoundError(f"对象不存在: {object_name}")
return data
def list_objects(self, prefix: str = "") -> list[str]:
"""List object names in the bucket that start with the given prefix."""
return self.client.list_objects(prefix=prefix)
def delete(self, object_name: str) -> None:
"""Handle delete for the Minio Document Binary Store instance."""
if not self.client.delete_object(object_name):
@@ -0,0 +1,142 @@
"""Postgres-backed persistence for cumulative AI model usage counters.
Keeps ModelUsageTracker (an in-memory, process-lifetime-only registry defined
in app/shared/model_usage_tracker.py) from losing its counters on every
backend restart. This store only ever persists the *current cumulative
snapshot* per provider+model not a historical time-series log matching
the "durable counters" scope decided in
docs/superpowers/specs/2026-07-23-status-model-usage-hardening-design.md.
"""
from __future__ import annotations
from contextlib import contextmanager
import psycopg2
import psycopg2.extras
from psycopg2.pool import ThreadedConnectionPool
from app.config.settings import settings
from app.shared.model_usage_tracker import ModelUsageEntry
# Table creation follows the same CREATE TABLE IF NOT EXISTS idiom used by
# every other Postgres store in this codebase — no migration framework.
_CREATE_TABLE = """
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)
);
"""
_UPSERT = """
INSERT INTO model_usage_stats
(provider, model, total_tokens, prompt_tokens, completion_tokens,
call_count_ok, call_count_error, last_called_at, last_latency_ms, last_error, updated_at)
VALUES
(%(provider)s, %(model)s, %(total_tokens)s, %(prompt_tokens)s, %(completion_tokens)s,
%(call_count_ok)s, %(call_count_error)s, %(last_called_at)s, %(last_latency_ms)s, %(last_error)s, NOW())
ON CONFLICT (provider, model) DO UPDATE SET
total_tokens = EXCLUDED.total_tokens,
prompt_tokens = EXCLUDED.prompt_tokens,
completion_tokens = EXCLUDED.completion_tokens,
call_count_ok = EXCLUDED.call_count_ok,
call_count_error = EXCLUDED.call_count_error,
last_called_at = EXCLUDED.last_called_at,
last_latency_ms = EXCLUDED.last_latency_ms,
last_error = EXCLUDED.last_error,
updated_at = NOW();
"""
class PostgresModelUsageStore:
"""Load and flush ModelUsageTracker snapshots to/from a Postgres table."""
def __init__(self) -> None:
"""Open a small connection pool and ensure the table exists."""
self._pool = ThreadedConnectionPool(
minconn=1,
maxconn=3,
host=settings.postgres_host,
port=settings.postgres_port,
user=settings.postgres_user,
password=settings.postgres_password,
dbname=settings.postgres_db,
)
self._ensure_schema()
def _ensure_schema(self) -> None:
"""Create the model_usage_stats table if it does not already exist."""
with self._conn() as conn:
with conn.cursor() as cur:
cur.execute(_CREATE_TABLE)
conn.commit()
@contextmanager
def _conn(self):
"""Borrow a pooled connection and always return it, even on error."""
conn = self._pool.getconn()
try:
yield conn
finally:
self._pool.putconn(conn)
def load_all(self) -> dict[str, ModelUsageEntry]:
"""Return every persisted row as {"provider:model": ModelUsageEntry}."""
with self._conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute("SELECT * FROM model_usage_stats")
rows = cur.fetchall()
entries: dict[str, ModelUsageEntry] = {}
for row in rows:
entry = ModelUsageEntry(
provider=row["provider"],
model=row["model"],
total_tokens=row["total_tokens"],
prompt_tokens=row["prompt_tokens"],
completion_tokens=row["completion_tokens"],
call_count_ok=row["call_count_ok"],
call_count_error=row["call_count_error"],
last_called_at=row["last_called_at"],
last_latency_ms=row["last_latency_ms"],
last_error=row["last_error"],
)
entries[f"{entry.provider}:{entry.model}"] = entry
return entries
def flush(self, entries: dict[str, ModelUsageEntry]) -> None:
"""Upsert the current cumulative snapshot of every tracked entry.
A no-op for an empty snapshot avoids opening a connection for nothing
(e.g. before any LLM/embedding/reranker call has happened yet).
"""
if not entries:
return
with self._conn() as conn:
with conn.cursor() as cur:
for entry in entries.values():
cur.execute(
_UPSERT,
{
"provider": entry.provider,
"model": entry.model,
"total_tokens": entry.total_tokens,
"prompt_tokens": entry.prompt_tokens,
"completion_tokens": entry.completion_tokens,
"call_count_ok": entry.call_count_ok,
"call_count_error": entry.call_count_error,
"last_called_at": entry.last_called_at,
"last_latency_ms": entry.last_latency_ms,
"last_error": entry.last_error,
},
)
conn.commit()
@@ -0,0 +1,5 @@
"""Celery task definitions for background processing.
This package exposes the shared Celery application instance and all
registered task functions used by API routes to enqueue work.
"""
@@ -0,0 +1,56 @@
"""Shared Celery application instance for background task processing.
All workers and enqueueing call sites import `celery_app` from this module
so the broker/backend configuration stays in one place.
"""
from __future__ import annotations
from celery import Celery
from app.config.settings import settings
def _redis_url() -> str:
"""Return a Redis connection URL from application settings."""
if settings.redis_password:
return (
f"redis://:{settings.redis_password}@"
f"{settings.redis_host}:{settings.redis_port}/{settings.redis_db}"
)
return f"redis://{settings.redis_host}:{settings.redis_port}/{settings.redis_db}"
_BROKER = _redis_url()
_BACKEND = _redis_url()
celery_app = Celery(
"compliance_hub",
broker=_BROKER,
backend=_BACKEND,
include=[
"app.infrastructure.tasks.document_tasks",
"app.infrastructure.tasks.perception_tasks",
],
)
celery_app.conf.update(
task_serializer="json",
result_serializer="json",
accept_content=["json"],
timezone="UTC",
enable_utc=True,
# Acknowledge task only after successful execution to avoid data loss.
task_acks_late=True,
task_reject_on_worker_lost=True,
# Keep results for 1 hour for status polling.
result_expires=3600,
# Scheduled counterpart to the Perception page's manual "Refresh" button.
# Only takes effect while a Beat process is running (./dev.sh start beat).
beat_schedule={
"crawl-regulations-periodic": {
"task": "app.infrastructure.tasks.perception_tasks.crawl_regulations_task",
"schedule": settings.perception_crawl_interval_seconds,
},
},
)
@@ -0,0 +1,73 @@
"""Celery tasks for document processing.
Each task is a thin wrapper that retrieves the already-stored document
binary and delegates to DocumentCommandService._process_document.
The task does not accept raw file bytes it reads them from the binary
store using the doc_id, so the Celery message payload stays small.
"""
from __future__ import annotations
from loguru import logger
from app.infrastructure.tasks.celery_app import celery_app
@celery_app.task(
name="app.infrastructure.tasks.document_tasks.process_document_task",
bind=True,
max_retries=3,
default_retry_delay=30,
acks_late=True,
)
def process_document_task(
self,
doc_id: str,
file_name: str,
doc_name: str,
regulation_type: str,
version: str,
generate_summary: bool,
run_id: str | None = None,
) -> dict:
"""Parse, embed, and index a document that has already been stored.
The task reads the file binary from MinIO using doc_id so the Celery
message stays small. Retries up to 3 times with a 30-second delay on
transient infrastructure errors.
"""
# Import inside the task function to avoid pickling issues and to ensure
# that each worker process initialises its own bootstrap singletons.
from app.shared.bootstrap import get_document_command_service, get_document_query_service
logger.info("process_document_task started: doc_id={}", doc_id)
try:
svc = get_document_command_service()
doc = get_document_query_service().get(doc_id)
if not doc:
raise ValueError(f"Document record not found: {doc_id}")
# Read the stored binary from MinIO — avoids passing raw bytes in the task message.
content = svc.binary_store.read(doc.object_name)
result = svc._process_document(
doc_id=doc_id,
file_name=file_name,
final_doc_name=doc_name,
content=content,
regulation_type=regulation_type,
version=version,
generate_summary=generate_summary,
run_id=run_id,
)
logger.info(
"process_document_task completed: doc_id={} status={} chunks={}",
doc_id, result.status, result.num_chunks,
)
return {"doc_id": result.doc_id, "status": result.status, "num_chunks": result.num_chunks}
except Exception as exc:
logger.exception("process_document_task failed: doc_id={}", doc_id)
# Retry on transient errors; permanent errors (bad file, parse failure)
# will exhaust retries and leave the document in FAILED state.
raise self.retry(exc=exc)
@@ -0,0 +1,63 @@
"""Celery task for scheduled regulatory source crawling.
This is the scheduled counterpart to the Perception page's manual "Refresh"
button (POST /perception/crawl). Every architecture reference document
describes source monitoring as continuous ("定时爬取"), not operator-triggered,
so this task is what Celery Beat runs on a fixed interval once an operator
starts a Beat process.
"""
from __future__ import annotations
from loguru import logger
from app.infrastructure.tasks.celery_app import celery_app
@celery_app.task(
name="app.infrastructure.tasks.perception_tasks.crawl_regulations_task",
bind=True,
)
def crawl_regulations_task(self) -> dict:
"""Crawl every registered regulatory source and enrich new/changed events.
Drains CrawlService.run_crawl(), which already isolates each source's
fetch and each event's enrichment behind its own try/except — a source
outage or a single bad event yields an "error" progress item and the
generator continues. Re-catching those here would only hide problems the
service has already handled, so this task's job is limited to counting
them and logging a summary.
No automatic retry is configured. An exception escaping run_crawl itself
means something broke in a way the service's own error handling did not
anticipate; the next scheduled tick already provides a retry within
settings.perception_crawl_interval_seconds, so an immediate retry against
the same failure is not worth the added complexity.
ponytail: relies on a single worker process to serialize scheduled runs
(Celery's default concurrency processes one task at a time, so a run that
outlasts the interval delays the next tick rather than overlapping it).
Add a Redis-based lock (e.g. SETNX on a per-task key) if this queue is
ever served by more than one worker.
"""
from app.shared.bootstrap import get_crawl_service
error_count = 0
new_count = 0
updated_count = 0
for item in get_crawl_service().run_crawl():
event = item.get("event")
if event == "error":
error_count += 1
logger.warning("Scheduled crawl source error: {}", item.get("data"))
elif event == "done":
data = item.get("data") or {}
new_count = data.get("total_new", 0)
updated_count = data.get("total_updated", 0)
logger.info(
"Scheduled crawl finished: new={} updated={} source_errors={}",
new_count, updated_count, error_count,
)
return {"new": new_count, "updated": updated_count, "source_errors": error_count}
@@ -9,6 +9,7 @@ from loguru import logger
from app.config.settings import settings
from app.domain.retrieval import Reranker, RetrievedChunk
from app.shared.model_usage_tracker import get_model_usage_tracker
class OpenAICompatibleReranker(Reranker):
@@ -37,10 +38,26 @@ class OpenAICompatibleReranker(Reranker):
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)],
@@ -54,22 +71,48 @@ class OpenAICompatibleReranker(Reranker):
return result
def _call_reranker(self, query: str, texts: list[str]) -> list[float]:
"""Call the reranker API and return a score per text."""
"""Call the reranker API and return a score per text.
Tries TEI format first (POST /rerank with model+texts), then falls back
to Cohere/OpenAI format (POST /v1/rerank with model+documents).
Both formats now include the model name, which most gateways require.
"""
headers = {"Content-Type": "application/json"}
if self._api_key:
headers["Authorization"] = f"Bearer {self._api_key}"
# Try TEI format first: POST /rerank
payload = {"query": query, "texts": texts, "raw_scores": False, "return_text": False}
# TEI format: POST /rerank — include model name (required by gateway proxies)
payload = {
"model": self._model,
"query": query,
"texts": texts,
"raw_scores": False,
"return_text": False,
}
url = f"{self._base_url}/rerank"
resp = requests.post(url, json=payload, headers=headers, timeout=self._timeout)
if resp.status_code == 404:
# Fall back to Cohere / OpenAI-style: POST /v1/rerank
if resp.status_code in (404, 400):
# Gateway returned an error — try Cohere/OpenAI-style format as fallback.
logger.debug(
"TEI rerank returned {} — trying Cohere format. Body: {}",
resp.status_code,
resp.text[:200],
)
payload_v1 = {"model": self._model, "query": query, "documents": texts}
url = f"{self._base_url}/v1/rerank"
resp = requests.post(url, json=payload_v1, headers=headers, timeout=self._timeout)
if not resp.ok:
# Surface a clear error message so callers can log it meaningfully.
try:
err_body = resp.json()
err_msg = err_body.get("error", {}).get("message", resp.text[:200])
except Exception:
err_msg = resp.text[:200]
resp.raise_for_status() # raises HTTPError with status code
raise ValueError(err_msg) # unreachable but satisfies type checker
resp.raise_for_status()
data = resp.json()
@@ -0,0 +1,21 @@
"""No-op reranker stub.
Returns the original candidate list sliced to top_k.
Replace with CrossEncoderReranker when a local cross-encoder model is available.
"""
from __future__ import annotations
from app.domain.retrieval.models import RetrievedChunk
from app.domain.retrieval.ports import Reranker
class PassThroughReranker(Reranker):
"""Pass-through reranker that preserves original retrieval order.
Acts as a placeholder for future cross-encoder reranking (e.g. ms-marco-MiniLM).
Wire via bootstrap.get_compliance_reranker() when ready to swap.
"""
def rerank(self, query: str, chunks: list[RetrievedChunk], top_k: int) -> list[RetrievedChunk]:
"""Return the first top_k chunks without reordering."""
return chunks[:top_k]
+8
View File
@@ -0,0 +1,8 @@
"""MCP (Model Context Protocol) server module.
Exposes selected read-only platform capabilities currently only regulation
search as MCP tools so external MCP clients (Claude Desktop, GitHub Copilot,
Cursor, etc.) can query this platform's compliance knowledge base directly.
"""
# Kept deliberately empty beyond this docstring — see server.py for the
# actual FastMCP instance and tool/middleware definitions.
+201
View File
@@ -0,0 +1,201 @@
"""MCPServer instance exposing the compliance knowledge base as an MCP tool.
This module is a pure protocol adapter: search_regulations() below calls the
existing AgentConversationService.ask() (the same application service backing
the /api/v1/agent/ask REST endpoint) and reshapes its result into a plain
dict. No new retrieval, ranking, or LLM orchestration logic lives here.
"""
from __future__ import annotations
import logging
import time
from typing import Annotated
from mcp.server import MCPServer
from mcp.server.transport_security import TransportSecuritySettings
from pydantic import Field
from starlette.responses import PlainTextResponse
from starlette.types import ASGIApp, Receive, Scope, Send
from app.config.settings import settings
from app.mcp.stats import get_mcp_stats_tracker
from app.shared.bootstrap import get_agent_conversation_service, get_jwt_handler
logger = logging.getLogger(__name__)
# Single shared MCPServer instance — analogous to the single shared FastAPI
# `app` instance in app/api/main.py. Tools registered via @mcp.tool() below.
# Note: the installed mcp SDK (2.0.0) renamed the older "FastMCP" class to
# "MCPServer" (mcp.server.mcpserver.MCPServer); the .tool()/.streamable_http_app()
# API surface used here is unchanged across that rename.
mcp = MCPServer("ai-regulations")
@mcp.tool()
def search_regulations(
query: Annotated[str, Field(min_length=1, max_length=2000)],
top_k: Annotated[int, Field(ge=1, le=20)] = 5,
) -> dict:
"""Search the compliance knowledge base and return a grounded answer.
query: Natural-language search question, e.g. "国六排放标准最新要求".
top_k: Maximum number of cited sources to return (1-20, default 5).
"""
# Bounds mirror AskRequest in app/api/models/agent.py so the MCP path cannot
# be used to bypass the REST endpoint's limits. They matter more here than
# there: KnowledgeRetrievalService amplifies top_k (candidate_k = top_k * 4)
# when reranking, and an LLM client can easily hallucinate a huge value.
# Declaring them via Annotated puts them in the advertised JSON schema too,
# so well-behaved clients never send an out-of-range value in the first place.
#
# No session_id is passed: this keeps each call stateless (no
# ConversationStore reads/writes), matching "search" semantics rather
# than multi-turn chat semantics.
started = time.perf_counter()
try:
_, result = get_agent_conversation_service().ask(query=query, top_k=top_k)
except Exception:
# Record the failure, then re-raise unchanged so the MCP SDK still
# converts it into a protocol-level error for the client. Swallowing
# it here would report success to the caller.
get_mcp_stats_tracker().record(
tool="search_regulations",
duration_ms=(time.perf_counter() - started) * 1000,
success=False,
)
raise
get_mcp_stats_tracker().record(
tool="search_regulations",
duration_ms=(time.perf_counter() - started) * 1000,
success=True,
)
return {
"answer": result.answer,
"sources": [source.__dict__ for source in result.sources],
}
class MCPAuthMiddleware:
"""Reject unauthenticated requests before they reach the MCP protocol handler.
Mirrors the existing get_current_user dependency's behavior (auth.py) but
implemented as raw ASGI middleware, since the mounted MCP app is a plain
ASGI app, not a FastAPI/APIRouter instance that supports Depends().
"""
def __init__(self, app: ASGIApp) -> None:
"""Store the wrapped ASGI app to delegate to once auth passes."""
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
"""Validate the bearer token for HTTP requests; pass non-HTTP scopes through."""
# Only HTTP requests carry an Authorization header to check; lifespan
# and other scope types must always pass through untouched.
if scope["type"] != "http" or not settings.auth_enabled:
await self.app(scope, receive, send)
return
headers = dict(scope["headers"])
# ASGI header values are raw bytes specified as latin-1, not UTF-8;
# decoding strictly as UTF-8 would raise on a malformed byte and turn a
# bad request into an unhandled 500.
auth_header = headers.get(b"authorization", b"").decode("latin-1")
token = auth_header.removeprefix("Bearer ").strip()
try:
get_jwt_handler().decode_token(token)
except ValueError as exc:
# Reject before the MCP session/protocol layer ever sees the request.
# WWW-Authenticate matches the get_current_user dependency (auth.py)
# and is required by RFC 7235 so clients can tell "needs credentials"
# apart from a generic failure.
response = PlainTextResponse(
str(exc), status_code=401, headers={"WWW-Authenticate": "Bearer"}
)
await response(scope, receive, send)
return
await self.app(scope, receive, send)
def _parse_allowed_hosts() -> list[str]:
"""Split the configured MCP host allow-list into individual entries."""
# Shared by the transport-security builder and the status endpoint so the
# panel can never display an allow-list different from the enforced one.
return [h.strip() for h in settings.mcp_allowed_hosts.split(",") if h.strip()]
def _build_transport_security() -> TransportSecuritySettings:
"""Translate the configured MCP host allow-list into SDK transport settings.
Without this the SDK infers its own allow-list from the bind host, which
defaults to 127.0.0.1 and therefore rejects every remote client with HTTP
421 fatal for a remotely deployed backend.
"""
allowed = _parse_allowed_hosts()
if "*" in allowed:
# Explicit, logged opt-out. Kept as an escape hatch for environments
# behind a proxy that rewrites Host unpredictably, but never the default.
logger.warning(
"MCP DNS-rebinding protection is disabled (mcp_allowed_hosts='*'). "
"Set MCP_ALLOWED_HOSTS to the real deployment host(s) instead."
)
return TransportSecuritySettings(enable_dns_rebinding_protection=False)
return TransportSecuritySettings(
enable_dns_rebinding_protection=True,
allowed_hosts=allowed,
# Browser clients send Origin; reuse the already-maintained CORS list so
# there is one place to declare trusted web origins. Non-browser MCP
# clients send no Origin at all, which the SDK treats as allowed.
allowed_origins=[o.strip() for o in settings.cors_allow_origins.split(",") if o.strip()],
)
def build_mcp_asgi_app() -> ASGIApp:
"""Return the Streamable HTTP ASGI app for the MCP server, auth-guarded.
streamable_http_path="/" is required here: MCPServer.streamable_http_app()
registers its own internal route at "/mcp" by default, and this app is
itself mounted at "/mcp" in api/main.py without overriding the internal
path to "/", the effective external path would be the confusing "/mcp/mcp"
instead of "/mcp".
"""
asgi_app = mcp.streamable_http_app(
streamable_http_path="/",
transport_security=_build_transport_security(),
)
asgi_app.add_middleware(MCPAuthMiddleware)
return asgi_app
async def get_mcp_status(public_url: str) -> dict:
"""Assemble the MCP status payload shown on the System Status page.
Owned by this module rather than the status route so that MCP internals
(the tool registry, the allow-list format, the stats tracker) stay behind
one boundary; the route only supplies public_url, which is the one value
only the HTTP layer can know.
"""
stats = get_mcp_stats_tracker().snapshot()
# list_tools() reads the in-memory registry populated by @mcp.tool() at
# import time, so the panel always reflects what is actually advertised
# rather than a hand-maintained duplicate list.
tools = await mcp.list_tools()
return {
"endpoint_url": public_url,
"auth_required": settings.auth_enabled,
"allowed_hosts": _parse_allowed_hosts(),
"tools": [
{
"name": tool.name,
"description": (tool.description or "").strip().split("\n")[0],
"calls": entry.calls if entry else 0,
"errors": entry.errors if entry else 0,
"avg_duration_ms": entry.avg_duration_ms if entry else None,
"last_called_at": (
entry.last_called_at.isoformat() if entry and entry.last_called_at else None
),
}
for tool, entry in ((tool, stats.get(tool.name)) for tool in tools)
],
}
+91
View File
@@ -0,0 +1,91 @@
"""In-memory per-tool call counters for the MCP server.
Lives in `app/mcp/` rather than `app/shared/` because these counters are
meaningful only for the MCP transport: they answer "is anything actually
calling our MCP endpoint, and does it work?" for the System Status page.
Token consumption is deliberately not tracked here MCP tool calls route
through AgentConversationService.ask() like every other caller, so the
existing ModelUsageTracker already accounts for it.
Counters are process-local and reset on restart. That is an accepted
tradeoff, recorded in the design spec: nothing billable depends on them.
"""
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 MCPToolStats:
"""Accumulated call outcomes for a single MCP tool."""
calls: int = 0
errors: int = 0
total_duration_ms: float = 0.0
last_called_at: datetime | None = None
@property
def avg_duration_ms(self) -> float | None:
"""Mean call duration, or None when the tool has never been called.
Returning None rather than 0.0 keeps "never called" distinguishable
from "called, but instantaneous" in the status UI.
"""
if self.calls == 0:
return None
return self.total_duration_ms / self.calls
class MCPStatsTracker:
"""Thread-safe registry of per-tool MCP call statistics.
The lock is load-bearing, not defensive habit: the mcp SDK dispatches
synchronous tool functions through anyio.to_thread.run_sync, so tool
bodies genuinely run on multiple worker threads at once unlike the
async REST routes, which are serialized by the event loop.
"""
def __init__(self) -> None:
"""Initialize an empty registry guarded by a single lock."""
self._tools: dict[str, MCPToolStats] = {}
# One coarse lock is enough: record() runs once per MCP tool call and
# snapshot() is only read by the low-traffic status endpoint.
self._lock = threading.Lock()
def record(self, *, tool: str, duration_ms: float, success: bool) -> None:
"""Record the outcome of one MCP tool invocation.
Never raises: a defect in observability code must not turn a working
tool call into a protocol error for the client.
"""
try:
# Coerce outside the lock so a bad argument cannot abort mid-update
# and leave calls incremented but duration unaccounted for.
duration = float(duration_ms)
now = datetime.now(timezone.utc)
with self._lock:
stats = self._tools.setdefault(tool, MCPToolStats())
stats.calls += 1
if not success:
stats.errors += 1
stats.total_duration_ms += duration
stats.last_called_at = now
except Exception as exc: # noqa: BLE001 - tracking must never break a real call
logger.warning("MCPStatsTracker.record failed for tool {} - {}", tool, exc)
def snapshot(self) -> dict[str, MCPToolStats]:
"""Return a shallow copy of all tracked tools, safe to read outside the lock."""
with self._lock:
return dict(self._tools)
@lru_cache
def get_mcp_stats_tracker() -> MCPStatsTracker:
"""Return the process-wide singleton tracker (mirrors get_model_usage_tracker())."""
return MCPStatsTracker()
+5
View File
@@ -12,6 +12,11 @@ class RagChatRequest(BaseModel):
top_k: int = 5
session_id: Optional[str] = None
filters: Optional[str] = None
# Optional document text to inject directly as LLM conversation context.
# When provided the document content is prepended to the query so the LLM
# can answer questions about it without requiring vector-store indexing.
context_text: Optional[str] = None
context_filename: Optional[str] = None
class RetrievedDoc(BaseModel):
+21 -2
View File
@@ -1,9 +1,16 @@
"""Provide service-layer logic for base client."""
"""Provide service-layer logic for base client.
P0-0: ``LLMResponse`` now carries an optional ``tool_calls`` list so that any
downstream code (agents, pipelines) can inspect and dispatch tool invocations
without touching the provider-specific adapter layer.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any
from enum import Enum
from app.services.llm.tool_types import Tool, ToolCall # noqa: F401 re-exported for callers
# Keep provider-specific behavior explicit so debugging stays straightforward.
@@ -24,6 +31,8 @@ class LLMResponse:
finish_reason: str = "stop"
latency_ms: int = 0
error: Optional[str] = None
# P0-0: populated when the model returns tool-call(s) instead of plain text.
tool_calls: List[ToolCall] = field(default_factory=list)
@property
def is_success(self) -> bool:
@@ -63,9 +72,19 @@ class BaseLLMClient(ABC):
messages: List[Dict[str, str]],
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List["Tool"]] = None,
**kwargs
) -> LLMResponse:
"""Handle chat for the Base L L M Client instance."""
"""Handle chat for the Base L L M Client instance.
Args:
messages: OpenAI-format message list.
max_tokens: Override config max_tokens when set.
temperature: Override config temperature when set.
tools: Optional list of Tool definitions to offer the model.
When provided, the model may respond with tool_calls in the
returned LLMResponse instead of (or in addition to) content.
"""
pass
def complete(
+50 -9
View File
@@ -1,11 +1,16 @@
"""Provide service-layer logic for deepseek client."""
"""Provide service-layer logic for deepseek client.
P0-0: ``chat()`` now accepts an optional ``tools`` list and parses ``tool_calls``
from the model response so that callers can dispatch tool invocations.
"""
import time
from typing import List, Dict, Optional
from typing import List, Dict, Optional, Generator
from loguru import logger
import httpx
from .base_client import BaseLLMClient, LLMResponse, LLMConfig, LLMProvider
from .tool_types import Tool, ToolCall
# Keep provider-specific behavior explicit so debugging stays straightforward.
@@ -46,13 +51,20 @@ class DeepSeekClient(BaseLLMClient):
messages: List[Dict[str, str]],
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Tool]] = None,
**kwargs
) -> LLMResponse:
"""Handle chat for the Deep Seek Client instance."""
"""Handle chat for the Deep Seek Client instance.
When ``tools`` is provided the request includes the tool definitions and
``tool_choice="auto"``; any tool_calls returned by the model are parsed
into ``LLMResponse.tool_calls``.
"""
import json
start_time = time.time()
try:
payload = {
payload: Dict = {
"model": self.config.model,
"messages": messages,
"max_tokens": max_tokens or self.config.max_tokens,
@@ -61,6 +73,11 @@ class DeepSeekClient(BaseLLMClient):
"stream": False
}
# P0-0: inject tool definitions when provided.
if tools:
payload["tools"] = [t.to_openai_format() for t in tools]
payload["tool_choice"] = "auto"
response = self._client.post("/chat/completions", json=payload)
response.raise_for_status()
@@ -71,12 +88,24 @@ class DeepSeekClient(BaseLLMClient):
choices = data.get("choices", [{}])
message = choices[0].get("message", {})
# P0-0: parse tool_calls returned by the model.
raw_tool_calls = message.get("tool_calls") or []
parsed_tool_calls: List[ToolCall] = []
for tc in raw_tool_calls:
fn = tc.get("function", {})
try:
args = json.loads(fn.get("arguments", "{}"))
except json.JSONDecodeError:
args = {}
parsed_tool_calls.append(ToolCall(id=tc.get("id", ""), name=fn.get("name", ""), arguments=args))
return LLMResponse(
content=message.get("content", ""),
content=message.get("content", "") or "",
model=data.get("model", self.config.model),
usage=data.get("usage", {}),
finish_reason=choices[0].get("finish_reason", "stop"),
latency_ms=latency_ms
latency_ms=latency_ms,
tool_calls=parsed_tool_calls,
)
except httpx.HTTPStatusError as e:
@@ -101,8 +130,14 @@ class DeepSeekClient(BaseLLMClient):
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
**kwargs
):
"""Stream chat for the Deep Seek Client instance."""
) -> Generator[str, None, Optional[Dict[str, int]]]:
"""Stream chat for the Deep Seek Client instance.
Returns the trailing token-usage dict as the generator's return value
(read via StopIteration.value when manually driven with next()) when
the gateway sends one via stream_options.include_usage, else None.
"""
usage: Optional[Dict[str, int]] = None
try:
payload = {
"model": self.config.model,
@@ -110,7 +145,8 @@ class DeepSeekClient(BaseLLMClient):
"max_tokens": max_tokens or self.config.max_tokens,
"temperature": temperature or self.config.temperature,
"top_p": kwargs.get("top_p", self.config.top_p),
"stream": True
"stream": True,
"stream_options": {"include_usage": True}
}
with self._client.stream("POST", "/chat/completions", json=payload) as response:
@@ -139,6 +175,9 @@ class DeepSeekClient(BaseLLMClient):
content = delta.get("content", "")
if content:
yield content
elif data.get("usage"):
# Trailing usage-only chunk — no content to yield, just capture it.
usage = data["usage"]
except json.JSONDecodeError:
continue
@@ -149,6 +188,8 @@ class DeepSeekClient(BaseLLMClient):
logger.error(f"DeepSeek Stream调用失败: {e}")
yield ""
return usage
def get_available_models(self) -> List[str]:
"""Return available models for the Deep Seek Client instance."""
return self.SUPPORTED_MODELS
+15 -6
View File
@@ -7,6 +7,8 @@ from functools import lru_cache
from .base_client import BaseLLMClient, LLMConfig, LLMProvider, LLMResponse
from .deepseek_client import DeepSeekClient
from .qwen_client import QwenClient, QwenVLClient
from .tracked_client import TrackedLLMClient
from app.shared.model_usage_tracker import get_model_usage_tracker
# Keep provider-specific behavior explicit so debugging stays straightforward.
@@ -14,7 +16,7 @@ from .qwen_client import QwenClient, QwenVLClient
# Keep provider-specific behavior explicit so debugging stays straightforward.
DEFAULT_MODELS = {
LLMProvider.DEEPSEEK: "deepseek-v4-flash",
LLMProvider.QWEN: "qwen3.5-flash",
LLMProvider.QWEN: "qwen3.6-flash",
LLMProvider.QWEN_VL: "qwen3-vl-plus"
}
@@ -45,7 +47,7 @@ class LLMFactory:
max_tokens: int = 4096,
temperature: float = 0.7,
**kwargs
) -> BaseLLMClient:
) -> "BaseLLMClient | TrackedLLMClient":
"""Handle create for the L L M Factory instance."""
provider_enum = self._parse_provider(provider)
@@ -76,11 +78,16 @@ class LLMFactory:
# 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] = client
LLMFactory._global_instances[cache_key] = tracked_client
logger.info(f"LLM客户端创建成功并缓存: {provider} - {model}")
return client
return tracked_client
def _parse_provider(self, provider: str) -> LLMProvider:
"""Handle parse provider for this module for the L L M Factory instance."""
@@ -94,6 +101,8 @@ class LLMFactory:
"qwen-max": LLMProvider.QWEN,
"qwen3.5-flash": LLMProvider.QWEN,
"qwen3.5-plus": LLMProvider.QWEN,
"qwen3.6-flash": LLMProvider.QWEN,
"qwen3.6-plus": LLMProvider.QWEN,
"qwen_vl": LLMProvider.QWEN_VL,
"qwen-vl": LLMProvider.QWEN_VL,
"qwen-vl-plus": LLMProvider.QWEN_VL,
@@ -137,7 +146,7 @@ class LLMFactory:
return client_class(config)
def get_cached(self, provider: str, model: Optional[str] = None) -> Optional[BaseLLMClient]:
def get_cached(self, provider: str, model: Optional[str] = None) -> "BaseLLMClient | TrackedLLMClient | None":
"""Return cached for the L L M Factory instance."""
provider_enum = self._parse_provider(provider)
model = model or DEFAULT_MODELS.get(provider_enum)
@@ -200,7 +209,7 @@ def get_llm_client(
provider: str = "qwen",
model: Optional[str] = None,
**kwargs
) -> BaseLLMClient:
) -> "BaseLLMClient | TrackedLLMClient":
"""Return llm client."""
factory = get_llm_factory()
+66 -12
View File
@@ -1,4 +1,8 @@
"""Provide service-layer logic for qwen client."""
"""Provide service-layer logic for qwen client.
P0-0: ``chat()`` now accepts an optional ``tools`` list and parses ``tool_calls``
from the model response so that callers can dispatch tool invocations.
"""
import time
import json
@@ -7,6 +11,7 @@ from loguru import logger
import httpx
from .base_client import BaseLLMClient, LLMResponse, LLMConfig, LLMProvider
from .tool_types import Tool, ToolCall
# Keep provider-specific behavior explicit so debugging stays straightforward.
@@ -22,6 +27,8 @@ class QwenClient(BaseLLMClient):
"qwen-long",
"qwen3.5-flash",
"qwen3.5-plus",
"qwen3.6-flash",
"qwen3.6-plus",
"qwen3-plus",
"qwen2.5-72b-instruct",
"qwen2.5-32b-instruct",
@@ -54,14 +61,20 @@ class QwenClient(BaseLLMClient):
messages: List[Dict[str, str]],
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
tools: Optional[List[Tool]] = None,
**kwargs
) -> LLMResponse:
"""Handle chat for the Qwen Client instance."""
"""Handle chat for the Qwen Client instance.
When ``tools`` is provided the request includes the tool definitions and
``tool_choice="auto"``; any tool_calls returned by the model are parsed
into ``LLMResponse.tool_calls``.
"""
start_time = time.time()
try:
# Keep provider-specific behavior explicit so debugging stays straightforward.
payload = {
payload: Dict = {
"model": self.config.model,
"messages": messages,
"max_tokens": max_tokens or self.config.max_tokens,
@@ -70,6 +83,11 @@ class QwenClient(BaseLLMClient):
"stream": False
}
# P0-0: inject tool definitions when provided.
if tools:
payload["tools"] = [t.to_openai_format() for t in tools]
payload["tool_choice"] = "auto"
# Keep provider-specific behavior explicit so debugging stays straightforward.
response = self._client.post("/chat/completions", json=payload)
response.raise_for_status()
@@ -82,12 +100,24 @@ class QwenClient(BaseLLMClient):
choices = data.get("choices", [{}])
message = choices[0].get("message", {})
# P0-0: parse tool_calls returned by the model.
raw_tool_calls = message.get("tool_calls") or []
parsed_tool_calls: List[ToolCall] = []
for tc in raw_tool_calls:
fn = tc.get("function", {})
try:
args = json.loads(fn.get("arguments", "{}"))
except json.JSONDecodeError:
args = {}
parsed_tool_calls.append(ToolCall(id=tc.get("id", ""), name=fn.get("name", ""), arguments=args))
return LLMResponse(
content=message.get("content", ""),
content=message.get("content", "") or "",
model=data.get("model", self.config.model),
usage=data.get("usage", {}),
finish_reason=choices[0].get("finish_reason", "stop"),
latency_ms=latency_ms
latency_ms=latency_ms,
tool_calls=parsed_tool_calls,
)
except httpx.HTTPStatusError as e:
@@ -112,8 +142,14 @@ class QwenClient(BaseLLMClient):
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
**kwargs
) -> Generator[str, None, None]:
"""Stream chat for the Qwen Client instance."""
) -> Generator[str, None, Optional[Dict[str, int]]]:
"""Stream chat for the Qwen Client instance.
Returns the trailing token-usage dict as the generator's return value
(read via StopIteration.value when manually driven with next()) when
the gateway sends one via stream_options.include_usage, else None.
"""
usage: Optional[Dict[str, int]] = None
try:
# Keep provider-specific behavior explicit so debugging stays straightforward.
payload = {
@@ -122,7 +158,8 @@ class QwenClient(BaseLLMClient):
"max_tokens": max_tokens or self.config.max_tokens,
"temperature": temperature or self.config.temperature,
"top_p": kwargs.get("top_p", self.config.top_p),
"stream": True # Keep provider-specific behavior explicit so debugging stays straightforward.
"stream": True, # Keep provider-specific behavior explicit so debugging stays straightforward.
"stream_options": {"include_usage": True}
}
# Keep provider-specific behavior explicit so debugging stays straightforward.
@@ -139,6 +176,9 @@ class QwenClient(BaseLLMClient):
data = json.loads(data_str)
choices = data.get("choices", [])
if not choices:
if data.get("usage"):
# Trailing usage-only chunk — capture it, nothing to yield.
usage = data["usage"]
continue # Keep provider-specific behavior explicit so debugging stays straightforward.
delta = choices[0].get("delta", {})
content = delta.get("content", "")
@@ -155,6 +195,8 @@ class QwenClient(BaseLLMClient):
logger.error(f"Qwen流式调用失败: {e}")
yield f"[ERROR: {str(e)}]"
return usage
async def async_stream_chat(
self,
messages: List[Dict[str, str]],
@@ -271,8 +313,14 @@ class QwenVLClient(BaseLLMClient):
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
**kwargs
) -> Generator[str, None, None]:
"""Stream chat for the Qwen V L Client instance."""
) -> Generator[str, None, Optional[Dict[str, int]]]:
"""Stream chat for the Qwen V L Client instance.
Returns the trailing token-usage dict as the generator's return value
(read via StopIteration.value when manually driven with next()) when
the gateway sends one via stream_options.include_usage, else None.
"""
usage: Optional[Dict[str, int]] = None
try:
payload = {
"model": self.config.model,
@@ -280,7 +328,8 @@ class QwenVLClient(BaseLLMClient):
"max_tokens": max_tokens or self.config.max_tokens,
"temperature": temperature or self.config.temperature,
"top_p": kwargs.get("top_p", self.config.top_p),
"stream": True
"stream": True,
"stream_options": {"include_usage": True}
}
with self._client.stream("POST", "/chat/completions", json=payload) as response:
@@ -295,6 +344,9 @@ class QwenVLClient(BaseLLMClient):
data = json.loads(data_str)
choices = data.get("choices", [])
if not choices:
if data.get("usage"):
# Trailing usage-only chunk — capture it, nothing to yield.
usage = data["usage"]
continue # Keep provider-specific behavior explicit so debugging stays straightforward.
delta = choices[0].get("delta", {})
content = delta.get("content", "")
@@ -307,6 +359,8 @@ class QwenVLClient(BaseLLMClient):
logger.error(f"QwenVL流式调用失败: {e}")
yield f"[ERROR: {str(e)}]"
return usage
def get_available_models(self) -> List[str]:
"""Return available models for the Qwen V L Client instance."""
return self.SUPPORTED_MODELS
@@ -319,7 +373,7 @@ class QwenVLClient(BaseLLMClient):
def create_qwen_client(
api_key: str,
model: str = "qwen3.5-flash",
model: str = "qwen3.6-flash",
base_url: str = "http://6.86.80.4:30080/v1",
**kwargs
) -> QwenClient:
+80
View File
@@ -0,0 +1,80 @@
"""Shared tool and tool-call type definitions for LLM function calling (P0-0).
These types implement the OpenAI-compatible tool/function-calling interface so that
any provider whose gateway supports the spec (DeepSeek, Qwen, etc.) can expose
tools to the LLM and receive structured tool invocations in return.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
@dataclass
class ToolCall:
"""Represent a single tool invocation returned by the LLM.
The model fills in ``id``, ``name``, and ``arguments`` when it decides to call
a tool instead of (or in addition to) producing a text response.
"""
# Unique identifier assigned by the model for this call.
id: str
# Name of the tool to invoke, matching the name registered in Tool.
name: str
# Parsed JSON arguments ready for direct use by the tool handler.
arguments: dict[str, Any] = field(default_factory=dict)
@dataclass
class ToolParameter:
"""JSON-Schemacompatible parameter block for a tool definition."""
# Top-level schema type — always "object" for OpenAI-compatible tools.
type: str = "object"
# Map of parameter name → JSON-Schema property descriptor.
properties: dict[str, Any] = field(default_factory=dict)
# List of required parameter names.
required: list[str] = field(default_factory=list)
@dataclass
class Tool:
"""Describe a callable tool that can be offered to the LLM.
Example usage::
search_tool = Tool(
name="search_regulations",
description="Search the compliance knowledge base for relevant regulation clauses.",
parameters=ToolParameter(
properties={"query": {"type": "string", "description": "Search query"}},
required=["query"],
),
)
response = client.chat(messages, tools=[search_tool])
"""
name: str
description: str
parameters: ToolParameter = field(default_factory=ToolParameter)
def to_openai_format(self) -> dict[str, Any]:
"""Serialise this tool to the OpenAI-compatible function-calling schema.
The returned dict can be placed directly in the ``tools`` list of a chat
completions request without any further transformation.
"""
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": {
"type": self.parameters.type,
"properties": self.parameters.properties,
"required": self.parameters.required,
},
},
}
@@ -0,0 +1,93 @@
"""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 and usage.
Drives the inner generator manually (instead of a plain `for` loop) so
it can capture the generator's return value via StopIteration.value —
the trailing token-usage dict the inner client captures from a
stream_options.include_usage chunk, if the gateway sent one.
"""
start = time.time()
error: Optional[str] = None
usage: Optional[Dict[str, int]] = None
gen = self._inner.stream_chat(messages, *args, **kwargs)
try:
while True:
try:
chunk = next(gen)
except StopIteration as stop:
usage = stop.value
break
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,
usage=usage,
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)
+221 -3
View File
@@ -2,10 +2,14 @@
from __future__ import annotations
import asyncio
from functools import lru_cache
from typing import Callable
from loguru import logger
from app.application.agent import AgentConversationService, AgentSessionService
from app.application.agent.agentic_service import AgenticConversationService
from app.application.documents import DocumentCommandService, DocumentQueryService
from app.application.knowledge import KnowledgeRetrievalService
from app.application.perception.services import PerceptionService
@@ -19,6 +23,17 @@ from app.infrastructure.parser.local_chunk_builder import LocalRegulationChunkBu
from app.infrastructure.parser.local_document_parser import LocalDocumentParser
from app.infrastructure.parser.vector_chunk_builder import AliyunVectorChunkBuilder
from app.infrastructure.perception.mock_event_store import MockEventStore
from app.infrastructure.perception.mock_notification_store import MockNotificationStore
from app.application.perception.crawl_service import CrawlService
from app.infrastructure.perception.base_event_store import BaseEventStore
from app.infrastructure.perception.base_notification_store import BaseNotificationStore
from app.infrastructure.perception.crawlers.catarc_crawler import CatarcCrawler
from app.infrastructure.perception.crawlers.guobiao_crawler import (
GuobiaoMandatoryCrawler,
GuobiaoRecommendedCrawler,
)
from app.infrastructure.perception.crawlers.eurlex_crawler import EurlexCrawler
from app.infrastructure.perception.llm_pipeline import LlmPipeline
from app.infrastructure.session.in_memory_conversation_store import InMemoryConversationStore
from app.infrastructure.storage.json_document_processing_store import JsonDocumentProcessingStore
from app.infrastructure.storage.json_document_repository import JsonDocumentRepository
@@ -26,11 +41,15 @@ from app.infrastructure.storage.minio_binary_store import MinioDocumentBinarySto
from app.infrastructure.storage.postgres_document_processing_store import PostgresDocumentProcessingStore
from app.infrastructure.storage.postgres_document_repository import PostgresDocumentRepository
from app.infrastructure.storage.postgres_parse_artifact_store import PostgresParseArtifactStore
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
from app.infrastructure.vectorstore.bm25_retriever import BM25Retriever
from app.infrastructure.vectorstore.cross_encoder_reranker import OpenAICompatibleReranker
from app.infrastructure.vectorstore.dense_retriever import DenseRetriever
from app.infrastructure.vectorstore.milvus_vector_index import MilvusVectorIndex
from app.services.llm.llm_factory import LLMFactory
from app.domain.compliance.ports import ComplianceRepository
from app.infrastructure.compliance.repository import PostgresComplianceRepository
from app.shared.model_usage_tracker import get_model_usage_tracker
# Keep shared wiring centralized so dependency construction remains consistent.
@@ -150,6 +169,14 @@ def get_parse_artifact_store():
return None
@lru_cache
def get_model_usage_store():
"""Return the Postgres model-usage store, or None when postgres backend is not enabled."""
if settings.document_repository_backend == "postgres":
return PostgresModelUsageStore()
return None
@lru_cache
def get_document_processing_store():
"""Return document processing store for the active repository backend."""
@@ -252,7 +279,31 @@ def get_document_query_service() -> DocumentQueryService:
@lru_cache
def get_conversation_store() -> InMemoryConversationStore:
"""Return conversation store."""
"""Return the active conversation store based on settings.
When session_backend='redis', sessions survive backend restarts and scale
across multiple API worker processes. When session_backend='memory' (default),
sessions are process-local and lost on restart.
"""
if settings.session_backend == "redis":
import redis as redis_lib
from app.infrastructure.session.redis_conversation_store import RedisConversationStore
# Build the Redis client from the same connection settings used by Celery.
kwargs: dict = {
"host": settings.redis_host,
"port": settings.redis_port,
"db": settings.redis_db,
"decode_responses": False,
}
if settings.redis_password:
kwargs["password"] = settings.redis_password
redis_client = redis_lib.Redis(**kwargs)
return RedisConversationStore( # type: ignore[return-value]
redis_client=redis_client,
timeout_seconds=settings.session_timeout_minutes * 60,
)
return InMemoryConversationStore(
max_sessions=settings.session_max_sessions,
timeout_minutes=settings.session_timeout_minutes,
@@ -269,26 +320,193 @@ def get_agent_conversation_service() -> AgentConversationService:
)
@lru_cache
def get_event_store() -> BaseEventStore:
"""Return event store selected by DOCUMENT_REPOSITORY_BACKEND setting."""
if settings.document_repository_backend == "postgres":
from app.infrastructure.perception.postgres_event_store import PostgresEventStore
return PostgresEventStore()
return MockEventStore()
@lru_cache
def get_notification_store() -> BaseNotificationStore:
"""Return notification store selected by DOCUMENT_REPOSITORY_BACKEND setting.
Mirrors get_event_store()'s gate: Mock in-memory when Postgres isn't
configured, so the feature works in local dev and tests without a
database.
"""
if settings.document_repository_backend == "postgres":
from app.infrastructure.perception.postgres_notification_store import (
PostgresNotificationStore,
)
return PostgresNotificationStore()
return MockNotificationStore()
@lru_cache
def get_compliance_repository() -> ComplianceRepository:
"""Return the compliance analysis repository.
Requires document_repository_backend=postgres and valid postgres_* settings.
Raises NotImplementedError for any other backend value.
"""
if settings.document_repository_backend != "postgres":
raise NotImplementedError(
f"ComplianceRepository requires document_repository_backend=postgres, "
f"got '{settings.document_repository_backend}'. "
"Set DOCUMENT_REPOSITORY_BACKEND=postgres in your .env file."
)
return PostgresComplianceRepository(
host=settings.postgres_host,
port=settings.postgres_port,
user=settings.postgres_user,
password=settings.postgres_password,
dbname=settings.postgres_db,
)
@lru_cache
def get_perception_service() -> PerceptionService:
"""Return perception service for regulatory intelligence."""
return PerceptionService(
event_store=MockEventStore(),
event_store=get_event_store(),
retrieval_service=get_retrieval_service(),
)
@lru_cache
def get_crawl_service() -> CrawlService:
crawlers = {
"CATARC": CatarcCrawler(),
"国标委·强制性": GuobiaoMandatoryCrawler(),
"国标委·推荐性": GuobiaoRecommendedCrawler(),
"EUR-Lex": EurlexCrawler(),
}
return CrawlService(
crawlers=crawlers,
event_store=get_event_store(),
llm_pipeline=LlmPipeline(),
retrieval_service=get_retrieval_service(),
notification_store=get_notification_store(),
embedding_provider=get_embedding_provider(),
vector_index=get_vector_index(),
)
@lru_cache
def get_agent_session_service() -> AgentSessionService:
"""Return agent session service."""
return AgentSessionService(conversation_store=get_conversation_store())
@lru_cache
def get_agentic_conversation_service() -> AgenticConversationService:
"""Return the Agentic RAG service (P0-1).
Uses the same retrieval, generation, and session infrastructure as the
standard chat service so no additional dependencies are required.
"""
return AgenticConversationService(
retrieval_service=get_retrieval_service(),
answer_generator=OpenAICompatibleAnswerGenerator(),
conversation_store=get_conversation_store(),
)
@lru_cache
def get_celery_app():
"""Return the shared Celery application instance.
Imported lazily so Celery is not required when running without workers
(e.g., tests that mock bootstrap or dev without Redis).
"""
from app.infrastructure.tasks.celery_app import celery_app
return celery_app
@lru_cache
def get_jwt_handler():
"""Return the shared JWTHandler instance for token creation and validation."""
from app.infrastructure.auth.jwt_handler import JWTHandler
return JWTHandler(
secret_key=settings.auth_secret_key,
algorithm=settings.auth_algorithm,
expire_minutes=settings.auth_token_expire_minutes,
)
@lru_cache
def get_user_store():
"""Return the PostgreSQL user store (lazy-connects on first call)."""
from app.infrastructure.auth.user_store import PostgresUserStore
return PostgresUserStore()
def preload_runtime_dependencies() -> None:
"""Warm dependencies that are safe and useful to preload during startup."""
LLMFactory.preload_clients(["qwen", "deepseek"])
_start_model_usage_persistence()
def cleanup_runtime_dependencies() -> None:
"""Release runtime dependencies that expose explicit cleanup hooks."""
LLMFactory.cleanup()
_stop_model_usage_persistence()
_model_usage_flush_task: "asyncio.Task | None" = None
def _start_model_usage_persistence() -> None:
"""Seed ModelUsageTracker from Postgres and start its periodic flush loop.
No-op when document_repository_backend != "postgres" ModelUsageTracker
then keeps behaving exactly as it always has: purely in-memory, reset on
every restart. Never raises: persistence must not block app startup.
"""
global _model_usage_flush_task
try:
store = get_model_usage_store()
except Exception as exc: # noqa: BLE001 - persistence must never block startup
logger.warning("Failed to initialize model usage persistence: {}", exc)
return
if store is None:
return
tracker = get_model_usage_tracker()
try:
tracker.seed(store.load_all())
except Exception as exc: # noqa: BLE001 - a bad load must not block startup
logger.warning("Failed to load persisted model usage stats: {}", exc)
async def _flush_loop() -> None:
"""Snapshot the tracker into Postgres every 60 seconds until cancelled."""
while True:
await asyncio.sleep(60)
try:
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)
_model_usage_flush_task = asyncio.create_task(_flush_loop())
def _stop_model_usage_persistence() -> None:
"""Cancel the periodic flush task and perform one best-effort final flush."""
global _model_usage_flush_task
if _model_usage_flush_task is not None:
_model_usage_flush_task.cancel()
_model_usage_flush_task = None
try:
store = get_model_usage_store()
except Exception as exc: # noqa: BLE001 - shutdown must not crash on this
logger.warning("Failed to access model usage store during shutdown: {}", exc)
return
if store is None:
return
try:
store.flush(get_model_usage_tracker().snapshot())
except Exception as exc: # noqa: BLE001 - shutdown must not crash on a flush failure
logger.warning("Failed final model usage flush: {}", exc)
+123
View File
@@ -0,0 +1,123 @@
"""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 seed(self, entries: dict[str, ModelUsageEntry]) -> None:
"""Bulk-load persisted entries (called once at startup, before any traffic).
Unlike record(), this replaces entries wholesale rather than
accumulating deltas it exists to restore counters saved by a
previous process run, not to record a new call.
"""
with self._lock:
self._entries.update(entries)
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()
+30 -3
View File
@@ -1,30 +1,57 @@
# ── Web framework ─────────────────────────────────────────────────────────────
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
python-multipart>=0.0.9
# MCP server module (backend/app/mcp/) — pin >=2.0.0: that release renamed the
# older "FastMCP" class to "MCPServer" (mcp.server.MCPServer), which is the
# class actually used in app/mcp/server.py.
mcp>=2.0.0
# ── Config & utilities ────────────────────────────────────────────────────────
pydantic>=2.0.0
pydantic-settings>=2.0.0
python-dotenv>=1.0.0
loguru>=0.7.0
httpx>=0.25.0
beautifulsoup4>=4.12.0
lxml>=5.0.0
tiktoken>=0.5.0
tenacity>=8.2.0
# Regulatory signal crawling (backend/app/infrastructure/perception/) — main-content
# extraction from crawled regulation detail pages and character-level diff for change
# detection. Import name for diff-match-patch is diff_match_patch (underscored).
trafilatura>=2.0.0
diff-match-patch>=20241021
# ── Auth ──────────────────────────────────────────────────────────────────────
python-jose[cryptography]>=3.3.0
# passlib is incompatible with bcrypt>=4.0 (removed __about__, strict 72-byte limit).
# Pin bcrypt to 3.x until passlib ships a fix.
passlib[bcrypt]>=1.7.4
bcrypt>=3.2.0,<4.0.0
# ── Async task queue ──────────────────────────────────────────────────────────
celery>=5.3.0
redis>=4.5.0
# ── Storage & databases ───────────────────────────────────────────────────────
pymilvus>=2.4.0
minio>=7.1.0
psycopg2-binary>=2.9.0
# ── Document parsing ─────────────────────────────────────────────────────────
pymupdf>=1.24.0
python-docx>=1.1.0
numpy>=1.24.0
alibabacloud-docmind-api20220711>=1.0.6
alibabacloud-tea-openapi>=0.3.11
alibabacloud-tea-util>=0.3.13
# ── RAG / LangChain ───────────────────────────────────────────────────────────
langchain>=0.1.0
langchain-milvus>=0.1.0
numpy>=1.24.0
# ── Testing ───────────────────────────────────────────────────────────────────
pytest>=7.4.0
pytest-asyncio>=0.21.0
fakeredis>=2.0.0
+140
View File
@@ -0,0 +1,140 @@
import asyncio
import pytest
from unittest.mock import MagicMock, patch
from datetime import datetime
from app.infrastructure.vectorstore.pass_through_reranker import PassThroughReranker
from app.domain.retrieval.models import RetrievedChunk
from app.domain.compliance.ports import AnalysisRecord, FindingRecord
# ── helpers ──────────────────────────────────────────────────────────────────
def _make_chunk(score: float) -> RetrievedChunk:
return RetrievedChunk(
chunk_id="c1",
doc_id="d1",
doc_title="Test Doc",
section_title="S1",
text="some text",
score=score,
page_start=1,
)
def _make_mock_client(content: str = '{"status":"ok","title":"T","desc":"D","clause_ref":"A1"}'):
client = MagicMock()
response = MagicMock()
response.is_success = True
response.content = content
client.chat.return_value = response
return client
def _make_mock_retrieval():
svc = MagicMock()
svc.retrieve.return_value = []
return svc
# ── existing tests ────────────────────────────────────────────────────────────
def test_pass_through_returns_top_k():
reranker = PassThroughReranker()
chunks = [_make_chunk(0.9), _make_chunk(0.8), _make_chunk(0.7)]
result = reranker.rerank(query="test", chunks=chunks, top_k=2)
assert len(result) == 2
assert result[0].score == 0.9
def test_pass_through_returns_all_when_top_k_exceeds():
reranker = PassThroughReranker()
chunks = [_make_chunk(0.5)]
result = reranker.rerank(query="test", chunks=chunks, top_k=10)
assert len(result) == 1
# ── new tests ─────────────────────────────────────────────────────────────────
def test_process_single_clause_returns_finding():
from app.application.compliance.pipeline import process_single_clause
client = _make_mock_client()
svc = _make_mock_retrieval()
result = process_single_clause("test clause", 0, svc, client)
assert result["finding"] is not None
assert result["index"] == 0
assert result["chunks"] == []
def test_run_clauses_parallel_runs_all():
from app.application.compliance.pipeline import run_clauses_parallel
client = _make_mock_client()
svc = _make_mock_retrieval()
clauses = ["clause one", "clause two", "clause three"]
results = asyncio.run(run_clauses_parallel(clauses, svc, client))
assert len(results) == 3
assert all(r["index"] == i for i, r in enumerate(results))
def test_run_clauses_parallel_handles_clause_failure():
from app.application.compliance.pipeline import run_clauses_parallel
svc = _make_mock_retrieval()
bad_client = MagicMock()
bad_client.chat.side_effect = RuntimeError("LLM exploded")
results = asyncio.run(run_clauses_parallel(
["clause one", "clause two"], svc, bad_client
))
assert len(results) == 2
assert all(r["finding"] is None for r in results)
assert all(r["chunks"] == [] for r in results)
# ── helpers for new tests ─────────────────────────────────────────────────────
def _sample_analysis() -> AnalysisRecord:
return AnalysisRecord(
id="a1", created_at=datetime(2026, 6, 8), created_by="u",
doc_name="doc.pdf", standard_name="EU AI Act",
risk_score=72, conclusion="Gaps found.", actions=[], para_text="para",
highlight_terms=[], findings=[],
)
def _sample_finding(status: str = "risk") -> FindingRecord:
return FindingRecord(
id="f1", analysis_id="a1", seq=0,
title="Missing CSMS", description="No CSMS certification.",
status=status, clause_ref="Art.9.1",
)
# ── new tests ─────────────────────────────────────────────────────────────────
def test_build_finding_context_contains_required_fields():
from app.application.compliance.pipeline import build_finding_context
ctx = build_finding_context(_sample_finding(), _sample_analysis())
assert "doc.pdf" in ctx
assert "EU AI Act" in ctx
assert "Missing CSMS" in ctx
assert "Art.9.1" in ctx
def test_generate_suggestions_returns_three_questions():
from app.application.compliance.pipeline import generate_suggestions
client = _make_mock_client(
'{"questions": ["Q1?", "Q2?", "Q3?"]}'
)
questions = generate_suggestions(_sample_finding("risk"), _sample_analysis(), client)
assert len(questions) == 3
assert all(isinstance(q, str) for q in questions)
def test_generate_suggestions_falls_back_on_error():
from app.application.compliance.pipeline import generate_suggestions
bad_client = MagicMock()
bad_resp = MagicMock()
bad_resp.is_success = False
bad_client.chat.return_value = bad_resp
questions = generate_suggestions(_sample_finding(), _sample_analysis(), bad_client)
assert len(questions) == 3 # fallback always returns 3
@@ -0,0 +1,98 @@
from unittest.mock import MagicMock, patch
from datetime import datetime
from app.domain.compliance.ports import (
AnalysisRecord,
FindingRecord,
ComplianceRepository,
)
def _mock_pool():
"""Return a mock psycopg2 ThreadedConnectionPool."""
conn = MagicMock()
cursor = MagicMock()
cursor.__enter__ = MagicMock(return_value=cursor)
cursor.__exit__ = MagicMock(return_value=False)
conn.cursor.return_value = cursor
pool = MagicMock()
pool.getconn.return_value = conn
return pool, conn, cursor
@patch("app.infrastructure.compliance.repository.psycopg2.pool.ThreadedConnectionPool")
def test_save_analysis_returns_uuid(mock_pool_cls):
from app.infrastructure.compliance.repository import PostgresComplianceRepository
pool, conn, cursor = _mock_pool()
mock_pool_cls.return_value = pool
cursor.fetchone.return_value = {"id": "abc-123"}
repo = PostgresComplianceRepository(
host="localhost", port=5432, user="u", password="p", dbname="db"
)
record = AnalysisRecord(
id="", created_at=datetime.utcnow(), created_by="user1",
doc_name="doc.pdf", standard_name="EU AI Act",
risk_score=50, conclusion="OK", actions=[], para_text="p",
highlight_terms=[], findings=[],
)
result = repo.save_analysis(record)
assert result == "abc-123"
def test_analysis_record_construction():
record = AnalysisRecord(
id="",
created_at=datetime.utcnow(),
created_by="user1",
doc_name="test.pdf",
standard_name="EU AI Act",
risk_score=72,
conclusion="Several gaps found.",
actions=[{"label": "Fix", "value": "Update docs"}],
para_text="The system shall...",
highlight_terms=["CSMS", "ISO 21434"],
findings=[
FindingRecord(
id="",
analysis_id="",
seq=0,
title="Missing CSMS",
description="No CSMS certification found.",
status="risk",
clause_ref="Art.9.1",
)
],
)
assert record.doc_name == "test.pdf"
assert len(record.findings) == 1
assert record.findings[0].status == "risk"
def test_compliance_repository_is_abstract():
import inspect
assert inspect.isabstract(ComplianceRepository)
def test_generate_docx_returns_bytes():
from app.infrastructure.compliance.docx_export import generate_docx
record = AnalysisRecord(
id="test-id", created_at=datetime(2026, 6, 8), created_by="user1",
doc_name="test.pdf", standard_name="EU AI Act",
risk_score=72, conclusion="Several gaps found.",
actions=[{"label": "Fix", "value": "Update CSMS docs"}],
para_text="The system shall implement CSMS.",
highlight_terms=["CSMS"],
findings=[
FindingRecord(
id="f1", analysis_id="test-id", seq=0,
title="Missing CSMS", description="No CSMS cert.",
status="risk", clause_ref="Art.9.1",
)
],
)
data = generate_docx(record)
assert isinstance(data, bytes)
assert len(data) > 1000 # DOCX is at minimum a ZIP with ~1 KB overhead
# Verify it's a valid ZIP (DOCX = ZIP container)
import zipfile, io
assert zipfile.is_zipfile(io.BytesIO(data))
+27
View File
@@ -0,0 +1,27 @@
"""Shared pytest fixtures and import-time guards for the backend test suite.
pytest imports this file before any test module beneath backend/tests/, which
makes it the only reliable place to install import-time guards: individual test
modules cannot guarantee they run first, because collection order follows
directory names.
"""
from __future__ import annotations
import sys
from unittest.mock import MagicMock
# app/shared/bootstrap.py (the composition root) eagerly imports the Postgres
# store modules, which do `import psycopg2` at their own module scope and later
# open a real connection pool. Any test that transitively imports bootstrap
# would therefore bind the real driver and attempt a live TCP connection to the
# configured production database, surfacing as a multi-second timeout rather
# than an obvious error. Binding mocks here — before the first test module is
# imported — makes that impossible regardless of collection order.
# setdefault (not assignment) keeps a real psycopg2 in place if something has
# already imported it deliberately.
_mock_psycopg2 = MagicMock()
_mock_psycopg2.extras = MagicMock()
sys.modules.setdefault("psycopg2", _mock_psycopg2)
sys.modules.setdefault("psycopg2.extras", _mock_psycopg2.extras)
sys.modules.setdefault("psycopg2.pool", MagicMock())
+2
View File
@@ -0,0 +1,2 @@
"""Test package for the MCP module (backend/app/mcp/)."""
# Empty package marker — no shared fixtures needed yet for this small test suite.
@@ -0,0 +1,93 @@
"""Unit tests for MCPAuthMiddleware.
Wraps a minimal dummy ASGI app (not the real MCP app) so these tests exercise
only the auth gate, not the MCP protocol itself keeps the test fast and
independent of FastMCP internals.
"""
from __future__ import annotations
from unittest.mock import patch
from starlette.applications import Starlette
from starlette.responses import PlainTextResponse
from starlette.routing import Route
from starlette.testclient import TestClient
from app.mcp.server import MCPAuthMiddleware
def _dummy_app() -> Starlette:
"""Build a minimal Starlette app that MCPAuthMiddleware can wrap."""
async def _ok(request):
"""Return a fixed 200 response so tests can assert pass-through."""
return PlainTextResponse("ok")
app = Starlette(routes=[Route("/ping", _ok)])
app.add_middleware(MCPAuthMiddleware)
return app
def test_missing_token_rejected_when_auth_enabled():
"""No Authorization header + auth_enabled=True -> 401."""
with patch("app.mcp.server.settings") as fake_settings:
fake_settings.auth_enabled = True
client = TestClient(_dummy_app())
response = client.get("/ping")
assert response.status_code == 401
def test_invalid_token_rejected_when_auth_enabled():
"""A token that fails decode_token() -> 401, request never reaches the app."""
fake_handler = type("H", (), {"decode_token": lambda self, t: (_ for _ in ()).throw(ValueError("bad token"))})()
with patch("app.mcp.server.settings") as fake_settings, \
patch("app.mcp.server.get_jwt_handler", return_value=fake_handler):
fake_settings.auth_enabled = True
client = TestClient(_dummy_app())
response = client.get("/ping", headers={"Authorization": "Bearer garbage"})
assert response.status_code == 401
def test_valid_token_passes_through_when_auth_enabled():
"""A token that decodes successfully -> request reaches the wrapped app."""
fake_handler = type("H", (), {"decode_token": lambda self, t: object()})()
with patch("app.mcp.server.settings") as fake_settings, \
patch("app.mcp.server.get_jwt_handler", return_value=fake_handler):
fake_settings.auth_enabled = True
client = TestClient(_dummy_app())
response = client.get("/ping", headers={"Authorization": "Bearer good"})
assert response.status_code == 200
assert response.text == "ok"
def test_auth_disabled_always_passes_through():
"""auth_enabled=False (dev mode) -> no token needed, matches get_current_user's dev bypass."""
with patch("app.mcp.server.settings") as fake_settings:
fake_settings.auth_enabled = False
client = TestClient(_dummy_app())
response = client.get("/ping")
assert response.status_code == 200
def test_401_includes_www_authenticate_header():
"""RFC 7235 requires WWW-Authenticate on 401 so clients can tell why they failed."""
with patch("app.mcp.server.settings") as fake_settings:
fake_settings.auth_enabled = True
client = TestClient(_dummy_app())
response = client.get("/ping")
assert response.status_code == 401
assert response.headers["WWW-Authenticate"] == "Bearer"
def test_non_utf8_authorization_header_is_rejected_not_crashed():
"""A non-UTF-8 header byte must yield a clean 401, not an unhandled 500.
ASGI header values are latin-1 bytes, so any remote client could otherwise
trigger a UnicodeDecodeError inside the middleware at will.
"""
with patch("app.mcp.server.settings") as fake_settings:
fake_settings.auth_enabled = True
client = TestClient(_dummy_app(), raise_server_exceptions=False)
# Bypass the http client's own header encoding by writing raw bytes.
response = client.get("/ping", headers={"Authorization": b"Bearer \xff\xfe"})
assert response.status_code == 401
+100
View File
@@ -0,0 +1,100 @@
"""Tests for the in-memory MCP per-tool statistics tracker.
These pin the two properties the status panel depends on: counters stay exact
under the concurrent thread dispatch the mcp SDK uses, and recording never
raises into a live tool call.
"""
from __future__ import annotations
import threading
from app.mcp.stats import MCPStatsTracker, MCPToolStats, get_mcp_stats_tracker
def test_avg_duration_is_none_before_any_call():
"""A never-called tool reports None, not 0.0, so the UI can distinguish them."""
assert MCPToolStats().avg_duration_ms is None
def test_avg_duration_is_the_mean_of_recorded_durations():
"""Average is computed over all calls, successful or not."""
tracker = MCPStatsTracker()
for duration in (100.0, 200.0, 300.0):
tracker.record(tool="search_regulations", duration_ms=duration, success=True)
stats = tracker.snapshot()["search_regulations"]
assert stats.calls == 3
assert stats.avg_duration_ms == 200.0
def test_failures_increment_both_calls_and_errors():
"""errors is a subset of calls, so the UI can show "2 of 3 failed" honestly."""
tracker = MCPStatsTracker()
tracker.record(tool="t", duration_ms=1.0, success=True)
tracker.record(tool="t", duration_ms=1.0, success=False)
tracker.record(tool="t", duration_ms=1.0, success=False)
stats = tracker.snapshot()["t"]
assert stats.calls == 3
assert stats.errors == 2
def test_last_called_at_is_set_and_timezone_aware():
"""The panel renders this as a local time, which requires an aware datetime."""
tracker = MCPStatsTracker()
tracker.record(tool="t", duration_ms=1.0, success=True)
last_called = tracker.snapshot()["t"].last_called_at
assert last_called is not None
assert last_called.tzinfo is not None
def test_concurrent_record_calls_are_not_lost():
"""8 threads x 100 calls must total exactly 800.
Without the lock this loses increments non-deterministically. The mcp SDK
runs synchronous tool bodies via anyio.to_thread.run_sync, so this is the
real dispatch model, not a hypothetical.
"""
tracker = MCPStatsTracker()
def hammer() -> None:
for _ in range(100):
tracker.record(tool="search_regulations", duration_ms=1.0, success=True)
threads = [threading.Thread(target=hammer) for _ in range(8)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
stats = tracker.snapshot()["search_regulations"]
assert stats.calls == 800
assert stats.total_duration_ms == 800.0
def test_record_swallows_bad_input_instead_of_raising():
"""A malformed duration must not propagate into the caller's tool call."""
tracker = MCPStatsTracker()
tracker.record(tool="t", duration_ms="not-a-number", success=True) # type: ignore[arg-type]
# Coercion happens before the lock is taken, so the entry is never created
# in a half-updated state.
assert tracker.snapshot() == {}
def test_snapshot_is_a_copy_not_the_live_dict():
"""Callers mutating the snapshot must not corrupt the tracker."""
tracker = MCPStatsTracker()
tracker.record(tool="t", duration_ms=1.0, success=True)
snapshot = tracker.snapshot()
snapshot.clear()
assert "t" in tracker.snapshot()
def test_get_mcp_stats_tracker_returns_a_singleton():
"""Instrumentation and the status route must observe the same counters."""
assert get_mcp_stats_tracker() is get_mcp_stats_tracker()
+79
View File
@@ -0,0 +1,79 @@
"""Tests for get_mcp_status(), the payload behind the System Status MCP card.
Covers the join between the live tool registry and the stats tracker, plus
the two values the route supplies or the settings decide.
"""
from __future__ import annotations
import asyncio
from unittest.mock import patch
from app.mcp.stats import MCPStatsTracker
def _status(tracker: MCPStatsTracker | None = None, **setting_overrides) -> dict:
"""Call get_mcp_status() with an isolated tracker and patched settings.
The real tracker is a process-wide singleton, so tests must inject their
own instance or they leak counters into each other.
"""
from app.mcp.server import get_mcp_status, settings
patched = settings.model_copy(update=setting_overrides)
with (
patch("app.mcp.server.settings", patched),
patch("app.mcp.server.get_mcp_stats_tracker", return_value=tracker or MCPStatsTracker()),
):
return asyncio.run(get_mcp_status("http://6.86.80.9:8000/mcp/"))
def test_public_url_is_passed_through_unmodified():
"""The route owns URL resolution; get_mcp_status() must not rewrite it."""
assert _status()["endpoint_url"] == "http://6.86.80.9:8000/mcp/"
def test_auth_required_follows_settings():
"""The panel's auth badge must reflect live config, not a hard-coded value."""
assert _status(auth_enabled=True)["auth_required"] is True
assert _status(auth_enabled=False)["auth_required"] is False
def test_allowed_hosts_are_split_and_stripped():
"""Displayed allow-list must match the one the transport actually enforces."""
status = _status(mcp_allowed_hosts="6.86.80.9:* , 127.0.0.1:*,")
assert status["allowed_hosts"] == ["6.86.80.9:*", "127.0.0.1:*"]
def test_tools_come_from_the_live_registry_with_zeroed_stats():
"""An advertised but never-called tool reports zeros, not absence."""
tools = {tool["name"]: tool for tool in _status()["tools"]}
assert "search_regulations" in tools
assert tools["search_regulations"]["calls"] == 0
assert tools["search_regulations"]["errors"] == 0
assert tools["search_regulations"]["avg_duration_ms"] is None
assert tools["search_regulations"]["last_called_at"] is None
def test_recorded_stats_are_joined_onto_the_matching_tool():
"""Counters recorded by the instrumented tool must surface on that tool's row."""
tracker = MCPStatsTracker()
tracker.record(tool="search_regulations", duration_ms=120.0, success=True)
tracker.record(tool="search_regulations", duration_ms=80.0, success=False)
tool = next(t for t in _status(tracker)["tools"] if t["name"] == "search_regulations")
assert tool["calls"] == 2
assert tool["errors"] == 1
assert tool["avg_duration_ms"] == 100.0
# Serialized for JSON transport; the frontend parses it with new Date().
assert isinstance(tool["last_called_at"], str)
def test_description_is_the_first_docstring_line():
"""Multi-line tool docstrings must not blow up the card's row height."""
tool = next(t for t in _status()["tools"] if t["name"] == "search_regulations")
assert "\n" not in tool["description"]
assert tool["description"].startswith("Search the compliance knowledge base")
@@ -0,0 +1,107 @@
"""Tests for the MCP endpoint's DNS-rebinding (Host header) protection.
The MCP SDK auto-enables DNS-rebinding protection and derives its allow-list
from the bind host, which defaults to 127.0.0.1. Left alone, that rejects every
request whose Host header is the real deployment address (6.86.80.9:8000) with
HTTP 421 before the auth middleware or the tool ever runs. These tests pin
the configured allow-list behavior so that failure mode cannot come back.
"""
from __future__ import annotations
import json
from contextlib import contextmanager
from unittest.mock import patch
from starlette.testclient import TestClient
from app.mcp.server import _build_transport_security, build_mcp_asgi_app
# A minimal JSON-RPC initialize call. Reaching the MCP handler at all is what
# matters here; transport security rejects the request long before this body is
# parsed, so its exact contents only need to be structurally valid.
_INITIALIZE = {
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2025-06-18",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"},
},
}
_HEADERS = {
"Content-Type": "application/json",
"Accept": "application/json, text/event-stream",
}
@contextmanager
def _mcp_client(allowed_hosts: str):
"""Yield a TestClient over the real MCP app with auth off and hosts configured.
The settings patch must stay active for the requests themselves, not just
for app construction, because MCPAuthMiddleware reads settings per request.
Entering the TestClient as a context manager is also required: it runs the
app's lifespan, without which the SDK's session manager task group is never
initialized and every request raises RuntimeError.
"""
with patch("app.mcp.server.settings") as fake_settings:
fake_settings.mcp_allowed_hosts = allowed_hosts
fake_settings.cors_allow_origins = "http://localhost:5173"
fake_settings.auth_enabled = False
with TestClient(build_mcp_asgi_app()) as client:
yield client
def test_remote_host_allowed_when_configured():
"""A configured non-loopback Host must reach the MCP handler, not 421."""
with _mcp_client("6.86.80.9:*,127.0.0.1:*") as client:
response = client.post(
"/", json=_INITIALIZE, headers={**_HEADERS, "Host": "6.86.80.9:8000"}
)
assert response.status_code == 200
assert "Invalid Host header" not in response.text
def test_unconfigured_host_still_rejected():
"""Protection must stay on: a Host outside the allow-list is refused with 421."""
with _mcp_client("6.86.80.9:*") as client:
response = client.post(
"/", json=_INITIALIZE, headers={**_HEADERS, "Host": "evil.example.com"}
)
assert response.status_code == 421
def test_initialize_response_is_event_stream():
"""Sanity check that a permitted request really completes the MCP handshake."""
with _mcp_client("6.86.80.9:*") as client:
response = client.post(
"/", json=_INITIALIZE, headers={**_HEADERS, "Host": "6.86.80.9:8000"}
)
assert response.status_code == 200
# The Streamable HTTP transport replies as SSE; the JSON-RPC result is
# embedded in a "data:" line rather than being the whole body.
payload = json.loads(response.text.split("data:", 1)[1].strip())
assert payload["result"]["serverInfo"]["name"] == "ai-regulations"
def test_wildcard_disables_protection_explicitly():
"""'*' is the documented opt-out; it must disable the check, not allow-list '*'."""
with patch("app.mcp.server.settings") as fake_settings:
fake_settings.mcp_allowed_hosts = "*"
fake_settings.cors_allow_origins = "http://localhost:5173"
security = _build_transport_security()
assert security.enable_dns_rebinding_protection is False
def test_allow_list_is_parsed_into_transport_settings():
"""Comma-separated config must become the SDK's allowed_hosts list verbatim."""
with patch("app.mcp.server.settings") as fake_settings:
fake_settings.mcp_allowed_hosts = "6.86.80.9:*, localhost:* ,"
fake_settings.cors_allow_origins = "http://localhost:5173"
security = _build_transport_security()
assert security.enable_dns_rebinding_protection is True
assert security.allowed_hosts == ["6.86.80.9:*", "localhost:*"]
assert security.allowed_origins == ["http://localhost:5173"]
@@ -0,0 +1,98 @@
"""Unit tests for the search_regulations MCP tool function.
Mocks AgentConversationService so no real retrieval/LLM call happens
verifies only the protocol-adapter contract: correct call shape in,
correct dict shape out.
"""
from __future__ import annotations
import asyncio
from dataclasses import dataclass
from unittest.mock import MagicMock, patch
@dataclass
class _FakeSource:
"""Minimal stand-in for a real Source dataclass (only __dict__ is used)."""
# A dataclass, not a MagicMock: the adapter serializes sources via
# source.__dict__, and a MagicMock's __dict__ is full of internal mock
# attributes, which would make the assertions meaningless.
doc_id: str
doc_title: str
score: float
@dataclass
class _FakeAnswerResult:
"""Minimal stand-in for AnswerResult — only .answer/.sources are read."""
answer: str
sources: list
def test_search_regulations_calls_agent_ask_without_session():
"""search_regulations must call ask() with no session_id (stateless search)."""
from app.mcp.server import search_regulations
fake_service = MagicMock()
fake_service.ask.return_value = (
None,
_FakeAnswerResult(answer="国六排放标准要求...", sources=[_FakeSource("doc-1", "国六标准", 0.9)]),
)
with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
search_regulations(query="国六排放标准最新要求", top_k=3)
fake_service.ask.assert_called_once_with(query="国六排放标准最新要求", top_k=3)
assert "session_id" not in fake_service.ask.call_args.kwargs
def test_search_regulations_shapes_response_dict():
"""The returned dict must expose 'answer' and 'sources' (list of plain dicts)."""
from app.mcp.server import search_regulations
fake_service = MagicMock()
fake_service.ask.return_value = (
None,
_FakeAnswerResult(answer="答案文本", sources=[_FakeSource("doc-2", "国标GB1589", 0.8)]),
)
with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
result = search_regulations(query="q")
assert result == {
"answer": "答案文本",
"sources": [{"doc_id": "doc-2", "doc_title": "国标GB1589", "score": 0.8}],
}
def test_search_regulations_default_top_k():
"""top_k defaults to 5 when the caller omits it."""
from app.mcp.server import search_regulations
fake_service = MagicMock()
fake_service.ask.return_value = (None, _FakeAnswerResult(answer="a", sources=[]))
with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
search_regulations(query="q")
assert fake_service.ask.call_args.kwargs["top_k"] == 5
def test_advertised_schema_bounds_top_k_and_query():
"""The advertised JSON schema must carry the same bounds as AskRequest.
Bounds declared via Annotated are what the SDK validates against and what
clients see, so asserting on the generated schema is the only way to catch
a regression that silently drops them.
"""
from app.mcp.server import mcp
schema = asyncio.run(mcp.list_tools())[0].input_schema["properties"]
assert schema["top_k"]["minimum"] == 1
assert schema["top_k"]["maximum"] == 20
assert schema["query"]["minLength"] == 1
assert schema["query"]["maxLength"] == 2000
@@ -0,0 +1,65 @@
"""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:<model>'."""
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.status_code = 500
failing_response.text = "boom"
failing_response.request = MagicMock()
failing_response.request.url = "http://example.com/embeddings"
failing_response.raise_for_status.side_effect = httpx.HTTPStatusError(
"boom", request=failing_response.request, response=failing_response
)
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
@@ -0,0 +1,44 @@
"""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
@@ -0,0 +1,104 @@
"""Unit tests for the model-usage persistence wiring in app.shared.bootstrap.
get_model_usage_store()'s settings-gating is tested the same way
tests/test_reranker_bootstrap.py tests get_reranker() by patching
"app.shared.bootstrap.settings" wholesale, matching this codebase's
established convention for testing @lru_cache settings-gated factories.
The remaining tests isolate _start_model_usage_persistence() /
_stop_model_usage_persistence() from get_model_usage_store() entirely (via
monkeypatch on the module-level function), so no real database or event loop
is needed anywhere in this file asyncio.create_task itself is also mocked.
"""
from __future__ import annotations
from unittest.mock import MagicMock, patch
# psycopg2 is mocked centrally in backend/tests/conftest.py, which pytest
# imports before any test module regardless of collection order.
from app.shared import bootstrap
from app.shared.model_usage_tracker import ModelUsageEntry, ModelUsageTracker
def test_get_model_usage_store_returns_none_when_not_postgres_backend():
"""get_model_usage_store() must be None unless document_repository_backend == 'postgres'."""
bootstrap.get_model_usage_store.cache_clear()
with patch("app.shared.bootstrap.settings") as mock_settings:
mock_settings.document_repository_backend = "json"
result = bootstrap.get_model_usage_store()
bootstrap.get_model_usage_store.cache_clear()
assert result is None
def test_get_model_usage_store_returns_instance_when_postgres_backend():
"""get_model_usage_store() must return a PostgresModelUsageStore when enabled.
ThreadedConnectionPool is mocked so no real connection is attempted; the
postgres_host/port/user/password/db values PostgresModelUsageStore reads
come from app.config.settings.settings directly (not from the
app.shared.bootstrap.settings reference mocked below), so they don't need
to be set here only document_repository_backend gates this factory.
"""
bootstrap.get_model_usage_store.cache_clear()
with patch("psycopg2.pool.ThreadedConnectionPool"), \
patch(
"app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema"
), \
patch("app.shared.bootstrap.settings") as mock_settings:
mock_settings.document_repository_backend = "postgres"
result = bootstrap.get_model_usage_store()
bootstrap.get_model_usage_store.cache_clear()
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
assert isinstance(result, PostgresModelUsageStore)
def test_start_model_usage_persistence_seeds_tracker_and_starts_flush_loop(monkeypatch):
"""When a store is available, startup must seed the tracker and schedule the flush task."""
fake_store = MagicMock()
fake_store.load_all.return_value = {
"deepseek:deepseek-v4-flash": ModelUsageEntry(
provider="deepseek", model="deepseek-v4-flash", total_tokens=99,
),
}
tracker = ModelUsageTracker()
monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: fake_store)
monkeypatch.setattr(bootstrap, "get_model_usage_tracker", lambda: tracker)
with patch("asyncio.create_task") as mock_create_task:
bootstrap._start_model_usage_persistence()
# Close the coroutine object passed to the mock so pytest doesn't warn
# about "coroutine was never awaited" — it was never meant to run here.
mock_create_task.call_args[0][0].close()
assert tracker.get("deepseek", "deepseek-v4-flash").total_tokens == 99
mock_create_task.assert_called_once()
bootstrap._stop_model_usage_persistence() # reset the module-level task handle
def test_start_model_usage_persistence_is_a_no_op_without_a_store(monkeypatch):
"""No store configured (json backend) — startup must not touch asyncio or the tracker."""
monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: None)
with patch("asyncio.create_task") as mock_create_task:
bootstrap._start_model_usage_persistence()
mock_create_task.assert_not_called()
def test_stop_model_usage_persistence_cancels_task_and_flushes(monkeypatch):
"""Shutdown must cancel the running flush task and perform one final flush."""
fake_store = MagicMock()
monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: fake_store)
fake_task = MagicMock()
bootstrap._model_usage_flush_task = fake_task
bootstrap._stop_model_usage_persistence()
fake_task.cancel.assert_called_once()
fake_store.flush.assert_called_once()
assert bootstrap._model_usage_flush_task is None
@@ -0,0 +1,102 @@
"""Unit tests for PostgresModelUsageStore, using a mocked psycopg2 pool.
Mirrors the mocking pattern in backend/tests/perception/test_postgres_event_store.py
no real database is needed.
"""
from __future__ import annotations
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
# psycopg2 is mocked centrally in backend/tests/conftest.py, so importing the
# module under test here never binds the real driver.
from app.shared.model_usage_tracker import ModelUsageEntry
def _cursor_returning(rows):
"""Build a MagicMock standing in for a psycopg2 cursor context manager."""
cursor = MagicMock()
cursor.__enter__ = lambda s: s
cursor.__exit__ = MagicMock(return_value=False)
cursor.fetchall.return_value = rows
return cursor
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
def test_load_all_returns_entries_keyed_by_provider_model(mock_pool_class, mock_ensure):
"""load_all() must turn each row into a ModelUsageEntry keyed by 'provider:model'."""
row = {
"provider": "deepseek",
"model": "deepseek-v4-flash",
"total_tokens": 100,
"prompt_tokens": 60,
"completion_tokens": 40,
"call_count_ok": 5,
"call_count_error": 1,
"last_called_at": datetime(2026, 7, 23, tzinfo=timezone.utc),
"last_latency_ms": 250,
"last_error": None,
}
mock_pool = MagicMock()
mock_pool_class.return_value = mock_pool
conn = MagicMock()
conn.__enter__ = lambda s: s
conn.__exit__ = MagicMock(return_value=False)
conn.cursor.return_value = _cursor_returning([row])
mock_pool.getconn.return_value = conn
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
store = PostgresModelUsageStore()
entries = store.load_all()
assert "deepseek:deepseek-v4-flash" in entries
entry = entries["deepseek:deepseek-v4-flash"]
assert isinstance(entry, ModelUsageEntry)
assert entry.total_tokens == 100
assert entry.call_count_error == 1
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
def test_flush_upserts_every_entry(mock_pool_class, mock_ensure):
"""flush() must execute one UPSERT per tracked entry and commit once."""
mock_pool = MagicMock()
mock_pool_class.return_value = mock_pool
conn = MagicMock()
conn.__enter__ = lambda s: s
conn.__exit__ = MagicMock(return_value=False)
cursor = MagicMock()
cursor.__enter__ = lambda s: s
cursor.__exit__ = MagicMock(return_value=False)
conn.cursor.return_value = cursor
mock_pool.getconn.return_value = conn
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
store = PostgresModelUsageStore()
entries = {
"deepseek:deepseek-v4-flash": ModelUsageEntry(
provider="deepseek", model="deepseek-v4-flash", total_tokens=100, call_count_ok=5,
),
}
store.flush(entries)
assert cursor.execute.call_count == 1
conn.commit.assert_called_once()
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
def test_flush_with_no_entries_does_not_touch_the_database(mock_pool_class, mock_ensure):
"""flush({}) must be a no-op — no point opening a connection for nothing."""
mock_pool = MagicMock()
mock_pool_class.return_value = mock_pool
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
store = PostgresModelUsageStore()
store.flush({})
mock_pool.getconn.assert_not_called()
@@ -0,0 +1,109 @@
"""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()
def test_seed_populates_registry_from_persisted_entries():
"""seed() must bulk-load entries (e.g. from Postgres at startup) into the registry."""
tracker = ModelUsageTracker()
persisted = {
"deepseek:deepseek-v4-flash": ModelUsageEntry(
provider="deepseek", model="deepseek-v4-flash", total_tokens=500, call_count_ok=20,
),
}
tracker.seed(persisted)
entry = tracker.get("deepseek", "deepseek-v4-flash")
assert entry.total_tokens == 500
assert entry.call_count_ok == 20
def test_seed_then_record_accumulates_on_top_of_seeded_value():
"""A call recorded after seeding must add to the seeded total, not replace it."""
tracker = ModelUsageTracker()
tracker.seed({
"deepseek:deepseek-v4-flash": ModelUsageEntry(
provider="deepseek", model="deepseek-v4-flash", total_tokens=500,
),
})
tracker.record(provider="deepseek", model="deepseek-v4-flash", success=True, usage={"total_tokens": 10})
assert tracker.get("deepseek", "deepseek-v4-flash").total_tokens == 510
@@ -0,0 +1,50 @@
"""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"
@@ -0,0 +1,116 @@
"""Unit tests verifying stream_chat() captures a trailing usage-only SSE chunk.
Exercises DeepSeekClient, QwenClient, and QwenVLClient directly (not through
TrackedLLMClient) by mocking the underlying httpx.Client.stream() call none
of these tests make a real network call.
"""
from __future__ import annotations
import json
from unittest.mock import MagicMock
from app.services.llm.base_client import LLMConfig, LLMProvider
from app.services.llm.deepseek_client import DeepSeekClient
def _sse_lines(*chunks: str, usage: dict | None = None) -> list[str]:
"""Build raw SSE 'data: ...' lines the way an OpenAI-compatible gateway sends them."""
lines = [
f'data: {json.dumps({"choices": [{"delta": {"content": c}}]})}'
for c in chunks
]
if usage is not None:
# Trailing usage-only chunk, as sent when stream_options.include_usage=true.
lines.append(f'data: {json.dumps({"choices": [], "usage": usage})}')
lines.append("data: [DONE]")
return lines
def _mock_streaming_client(lines: list[str]) -> MagicMock:
"""Build a MagicMock standing in for httpx.Client, configured for .stream()."""
fake_response = MagicMock()
fake_response.raise_for_status.return_value = None
fake_response.iter_lines.return_value = lines
stream_cm = MagicMock()
stream_cm.__enter__.return_value = fake_response
stream_cm.__exit__.return_value = False
client = MagicMock()
client.stream.return_value = stream_cm
return client
def _drain(gen):
"""Manually drive a generator, returning (yielded_chunks, stop_iteration_value)."""
chunks = []
value = None
while True:
try:
chunks.append(next(gen))
except StopIteration as stop:
value = stop.value
break
return chunks, value
def test_deepseek_stream_chat_returns_usage_from_trailing_chunk():
"""DeepSeekClient.stream_chat() must return the trailing usage dict."""
config = LLMConfig(provider=LLMProvider.DEEPSEEK, model="deepseek-v4-flash", api_key="k", base_url="http://x/v1")
client = DeepSeekClient(config)
usage = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
client._client = _mock_streaming_client(_sse_lines("Hello", " world", usage=usage))
chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "hi"}]))
assert chunks == ["Hello", " world"]
assert returned_usage == usage
# The gateway must actually be asked to include usage in the stream.
sent_payload = client._client.stream.call_args.kwargs["json"]
assert sent_payload["stream_options"] == {"include_usage": True}
def test_deepseek_stream_chat_without_usage_chunk_returns_none():
"""If the gateway never sends a usage chunk, the generator returns None (unchanged behavior)."""
config = LLMConfig(provider=LLMProvider.DEEPSEEK, model="deepseek-v4-flash", api_key="k", base_url="http://x/v1")
client = DeepSeekClient(config)
client._client = _mock_streaming_client(_sse_lines("Hi"))
chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "hi"}]))
assert chunks == ["Hi"]
assert returned_usage is None
from app.services.llm.qwen_client import QwenClient, QwenVLClient
def test_qwen_stream_chat_returns_usage_from_trailing_chunk():
"""QwenClient.stream_chat() must return the trailing usage dict."""
config = LLMConfig(provider=LLMProvider.QWEN, model="qwen3.5-flash", api_key="k", base_url="http://x/v1")
client = QwenClient(config)
usage = {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}
client._client = _mock_streaming_client(_sse_lines("Bonjour", usage=usage))
chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "hi"}]))
assert chunks == ["Bonjour"]
assert returned_usage == usage
sent_payload = client._client.stream.call_args.kwargs["json"]
assert sent_payload["stream_options"] == {"include_usage": True}
def test_qwen_vl_stream_chat_returns_usage_from_trailing_chunk():
"""QwenVLClient.stream_chat() must return the trailing usage dict."""
config = LLMConfig(provider=LLMProvider.QWEN_VL, model="qwen3-vl-plus", api_key="k", base_url="http://x/v1")
client = QwenVLClient(config)
usage = {"prompt_tokens": 20, "completion_tokens": 6, "total_tokens": 26}
client._client = _mock_streaming_client(_sse_lines("Describing image", usage=usage))
chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "describe"}]))
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}
@@ -0,0 +1,105 @@
"""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."""
# Use MagicMock to avoid requiring a real LLM provider implementation (e.g., DeepseekClient);
# tests focus on TrackedLLMClient's delegation and recording behavior, not provider logic.
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
def test_stream_chat_records_usage_from_generator_return_value():
"""stream_chat() must forward the inner generator's returned usage dict to record()."""
inner = _make_inner()
def fake_stream(*args, **kwargs):
yield "chunk-1"
yield "chunk-2"
return {"prompt_tokens": 6, "completion_tokens": 2, "total_tokens": 8}
inner.stream_chat.side_effect = fake_stream
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.total_tokens == 8
assert entry.call_count_ok == 1

Some files were not shown because too many files have changed in this diff Show More