From 9a9dbd7cd25c92e333fa3ec92f4babd783a4aee3 Mon Sep 17 00:00:00 2001 From: Rookiecoder-jsjs Date: Tue, 28 Jul 2026 22:18:21 +0800 Subject: [PATCH] feat: deep-think toggle, i18n, followup removal, storage hardening MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Chat UX: - Deep Think toggle (Qwen hybrid-thinking): reasoning_content streamed as `event: thinking` SSE frames before the answer body; 400-fallback retries without enable_thinking for model tiers that reject the param; frontend toggle persists in localStorage, collapsible reasoning block with timing. - Follow-up question chips removed end-to-end (model field, API path, SSE done-frame, frontend chips/handler/css, tests, eval gold case). - ChatPage UI translated to Chinese (timing chip, sources, feedback, errors). Storage / ops: - docker-compose: bypass chromadb 0.4.18 entrypoint's `pip install --force-reinstall chroma-hnswlib` (pulls numpy 2.x → np.float_ removed → crash loop); run uvicorn directly. Real bind-mount persistence to ./data/chromadb; ./data/neo4j bind-mounts for graph data. - scripts/rebuild_chroma.py: zero-API-cost vector rebuild from SQLite chunks + embedding_cache (md5-keyed). Reuses get_chroma_client so collection name/cosine/upsert match ingestion; metadata fields mirror documents.py:287-296 verbatim. Recovery path for future vector loss. Infra (earlier batch): - SSE client util + doc_status service + rate_limit middleware. - Embedding cache self-heal (corrupt JSON/pickle blobs pruned). - Neo4j batch chunk/link writes (UNWIND) replacing per-chunk round-trips. - entity_extractor refactor; progress_tracker; reranker. - tests: +conftest, +test_doc_status, +test_sse; -test_followups, -eval/06. Verified: 54 chunks repopulated into Chroma (0 cache misses, 0 API calls); collection count=54; metadata user_id/document_id/prev_next intact. --- .gitignore | 3 + CLAUDE.md | 155 +++++++ README.md | 4 +- backend/app/api/auth.py | 36 +- backend/app/api/chat.py | 206 ++++++--- backend/app/api/documents.py | 164 +++++-- backend/app/api/progress.py | 46 +- backend/app/auth/rate_limit.py | 57 +++ backend/app/config.py | 57 ++- backend/app/database.py | 44 +- backend/app/main.py | 10 +- backend/app/models/chat.py | 10 +- backend/app/models/document.py | 11 + backend/app/models/user.py | 10 +- backend/app/services/chroma_client.py | 5 +- backend/app/services/doc_status.py | 202 +++++++++ backend/app/services/embedding.py | 115 +++-- backend/app/services/entity_extractor.py | 280 ++++-------- backend/app/services/llm.py | 400 ++++++++++++----- backend/app/services/neo4j_client.py | 98 +++- backend/app/services/progress_tracker.py | 50 ++- backend/app/services/reranker.py | 8 + backend/eval/gold/06_followup.json | 27 -- backend/scripts/rebuild_chroma.py | 139 ++++++ backend/tests/conftest.py | 15 + backend/tests/test_doc_status.py | 219 +++++++++ backend/tests/test_document_detail.py | 2 + backend/tests/test_embedding.py | 1 + backend/tests/test_followups.py | 358 --------------- docker-compose.yml | 21 +- frontend/src/api/chat.js | 10 +- frontend/src/utils/sse.js | 62 +++ frontend/src/views/ChatPage.vue | 546 ++++++++++++++++------- frontend/src/views/DocumentsPage.vue | 2 +- frontend/tests/test_sse.cjs | 130 ++++++ 35 files changed, 2430 insertions(+), 1073 deletions(-) create mode 100644 CLAUDE.md create mode 100644 backend/app/auth/rate_limit.py create mode 100644 backend/app/services/doc_status.py delete mode 100644 backend/eval/gold/06_followup.json create mode 100644 backend/scripts/rebuild_chroma.py create mode 100644 backend/tests/conftest.py create mode 100644 backend/tests/test_doc_status.py delete mode 100644 backend/tests/test_followups.py create mode 100644 frontend/src/utils/sse.js create mode 100644 frontend/tests/test_sse.cjs diff --git a/.gitignore b/.gitignore index 7a5b9b5..8b4f5a9 100644 --- a/.gitignore +++ b/.gitignore @@ -49,3 +49,6 @@ RAG_检索优化指南.md # Local IDE / agent config .claude/ .trae/ + +# CodeGraph index (regenerable tool artifact) +.codegraph/ diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..035dbbd --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,155 @@ +# CLAUDE.md — Knowledge Graph System + +## CodeGraph 使用规范 + +本项目已接入 CodeGraph MCP，遵循以下规则： + +### 初次加载 / 全新仓库 +```bash +# 克隆或首次打开项目时，必须先初始化索引 +codegraph init +``` + +### 代码修改后 +- **自动同步**：文件保存后 CodeGraph 监听器会在约 2 秒内自动更新索引，无需手动操作 +- **批量修改后**（如 git pull、大规模重构）：运行一次手动同步确保一致 + ```bash + codegraph sync + ``` +- **验证索引状态**： + ```bash + codegraph status + ``` + +### 回答代码问题时 +- 优先使用 `codegraph_explore` 工具，**不要**逐文件 grep / read +- 典型查询：调用链、模块依赖、改动影响范围、"X 是怎么工作的" +- CodeGraph 返回的代码片段视为**已读取**，无需再次打开文件 +- 若索引中有 `⚠️` 过期标记，再用 Read 工具读取对应文件的最新内容 + +### CLI 速查 +```bash +codegraph query # 搜索符号 +codegraph callers # 谁调用了它 +codegraph callees # 它调用了谁 +codegraph impact # 改动影响范围 +codegraph explore # 自然语言探索（等同 MCP 工具） +``` + +--- + +## 项目概览 + +**Knowledge Graph System** — 多用户知识图谱平台，支持文档上传、实体抽取、语义搜索和 AI 问答。 + +| 层 | 技术栈 | +|----|--------| +| 后端 | FastAPI + Python 3.x | +| 图数据库 | Neo4j（实体/关系存储）| +| 向量数据库 | ChromaDB（Embedding 检索）| +| 关系数据库 | SQLite（用户/文档元数据）| +| 前端 | Vue 3 + Vite | +| LLM | 百炼 Qwen / Kimi / SiliconFlow（可切换）| +| Embedding | Qwen3-Embedding-8B（via SiliconFlow）| +| Reranker | Qwen3-Reranker-8B | + +--- + +## 目录结构 + +``` +D:\NC/ +├── backend/ +│ ├── app/ +│ │ ├── main.py # FastAPI 入口，lifespan 管理 Neo4j/Chroma 连接 +│ │ ├── config.py # pydantic-settings，所有环境变量集中管理 +│ │ ├── database.py # SQLite 初始化 +│ │ ├── logger.py # 日志配置 +│ │ ├── api/ # 路由层（auth, documents, search, graph, chat, +│ │ │ # progress, tags, timeline, dashboard） +│ │ ├── models/ # Pydantic 数据模型 +│ │ ├── services/ # 业务逻辑层 +│ │ │ ├── neo4j_client.py # Neo4j 图操作 +│ │ │ ├── chroma_client.py # ChromaDB 向量检索 +│ │ │ ├── embedding.py # 向量化 +│ │ │ ├── entity_extractor.py # LLM 实体抽取 +│ │ │ ├── chunker.py # 文档分块 +│ │ │ ├── bm25.py # 关键词检索 +│ │ │ ├── fusion.py # BM25 + 向量融合排序 +│ │ │ ├── reranker.py # 重排序 +│ │ │ ├── llm.py # LLM 调用封装 +│ │ │ └── query_processor.py # 查询处理 +│ │ └── auth/ # JWT 认证（security.py, jwt_handler.py） +│ ├── tests/ # pytest 测试 +│ └── eval/ # 检索质量评估 +├── frontend/ +│ └── src/ +│ ├── views/ # 页面（Dashboard, Documents, Chat, Graph, +│ │ # Search, Timeline, EntityDetail, ClusterMap） +│ └── components/ # UI 组件 + 布局 +└── .codegraph/ # CodeGraph 索引（勿手动修改） +``` + +--- + +## 开发环境 + +### 必需服务 +| 服务 | 默认地址 | 说明 | +|------|----------|------| +| Neo4j | bolt://localhost:7687 | 图数据库 | +| ChromaDB | localhost:8000 | 向量数据库 | + +### 环境变量（`.env`） +```bash +# 必填 — 启动时会校验，占位符会导致 RuntimeError +JWT_SECRET=<用 python -c "import secrets; print(secrets.token_urlsafe(48))" 生成> + +# LLM（至少配置一个） +BAILIAN_API_KEY=... +# SILICON_FLOW_API_KEY=... +# KIMI_API_KEY=... + +# 数据库（有默认值，按需覆盖） +NEO4J_URI=bolt://localhost:7687 +NEO4J_USER=neo4j +NEO4J_PASSWORD=... +CHROMA_HOST=localhost +CHROMA_PORT=8000 +``` + +### 启动 +```bash +# 后端 +cd backend +.venv/Scripts/activate # Windows +uvicorn app.main:app --reload --port 8001 + +# 前端 +cd frontend +npm install +npm run dev # 默认 http://localhost:5173 +``` + +--- + +## 测试 + +```bash +cd backend +pytest tests/ -v +pytest tests/ --cov=app --cov-report=term-missing # 带覆盖率 +``` + +测试文件与功能模块对应：`test_graph_rag.py`, `test_embedding.py`, `test_tags.py`, `test_timeline.py`, `test_dashboard.py` 等。 + +--- + +## 编码约定 + +- **Python**：PEP 8，所有函数加类型注解，用 `logging` 不用 `print` +- **错误处理**：API 层统一返回结构化错误，服务层抛出有意义的异常，不吞掉错误 +- **不可变优先**：数据对象用 `@dataclass(frozen=True)` 或 Pydantic model，避免原地修改 +- **配置集中**：所有配置项在 `app/config.py`，不在业务代码里读 `os.environ` +- **安全**：密钥只走环境变量，JWT_SECRET 禁止使用占位符，CORS 不用 `*` +- **文件体量**：单文件不超过 400 行，超出则拆分为子模块 diff --git a/README.md b/README.md index 073437a..edaa0ba 100644 --- a/README.md +++ b/README.md @@ -202,7 +202,7 @@ Vite 已配置 `/api` 代理到 `http://localhost:8001`。 ## 🔌 API 接口设计 -所有需要鉴权的接口都要求 `Authorization: Bearer ` header。SSE 进度接口也通过 header 鉴权，**不接受** URL `?token=` 参数。 +所有需要鉴权的接口都要求 `Authorization: Bearer ` header。SSE 进度接口**优先**使用 header 鉴权；但由于原生 `EventSource` 客户端无法设置自定义 header，作为兼容回退也接受 `?token=` 查询参数。⚠️ 查询参数形式会把 token 泄露进反向代理访问日志与浏览器历史，请将此类 URL 视为敏感信息（前端 `EventSource` 即使用此回退方式）。 ### 🔐 认证 `/api/auth` - `POST /api/auth/register` — 用户注册（用户名 3-50 字符 `[A-Za-z0-9_.-]`、密码 ≥ 8 字符且必须含字母+数字） @@ -325,7 +325,7 @@ PDF/Word/TXT/MD → markitdown → Markdown → 层级解析 → 语义切块 | 🗝️ API 密钥 | 环境变量 | **勿**硬编码到代码；`.env` 已 gitignore | | 🚦 生产模式 | `APP_ENV=production` | 默认 JWT_SECRET 启动时直接 `RuntimeError` | | 🚪 401 处理 | 拦截器去重 | 防重入 + 派发 `auth:logout` 事件 | -| 📡 进度 SSE | Authorization header | 不接受 `?token=` URL 参数（避免日志泄露） | +| 📡 进度 SSE | Authorization header 优先 | EventSource 回退接受 `?token=`（会进代理日志，视为敏感 URL） | | 💾 嵌入缓存 | JSON 序列化 | 取代 `pickle`（防反序列化漏洞） | | ✅ 注册校验 | 强校验 | 用户名 `[A-Za-z0-9_.-]`、密码 ≥ 8 字符含字母+数字 | | 👥 Neo4j 删除 | 跨用户隔离 | `delete_document` step 4 强制 `user_id` 过滤 | diff --git a/backend/app/api/auth.py b/backend/app/api/auth.py index 629b967..2e59285 100644 --- a/backend/app/api/auth.py +++ b/backend/app/api/auth.py @@ -2,7 +2,7 @@ from datetime import timedelta from typing import Optional -from fastapi import APIRouter, Depends, HTTPException, status +from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm import aiosqlite @@ -10,6 +10,7 @@ from app.database import get_db from app.auth.security import verify_password, get_password_hash from app.auth.jwt_handler import create_access_token, verify_token +from app.auth.rate_limit import login_limiter, register_limiter from app.models.user import UserCreate, UserResponse, Token router = APIRouter(prefix="/api/auth", tags=["auth"]) @@ -17,6 +18,26 @@ oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/auth/login") +def _client_ip(request: Optional[Request]) -> str: + """Best-effort client IP; 'unknown' for direct (non-HTTP) calls.""" + if request is not None and request.client is not None: + return request.client.host + return "unknown" + + +def _rate_limit(limiter, key: str) -> None: + """Raise 429 when ``key`` exceeds ``limiter``. Skipped under test.""" + settings = get_settings() + if settings.APP_ENV.lower() in ("test", "testing"): + return # never throttle the automated test suite + if not limiter.is_allowed(key): + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail="Too many attempts. Please try again later.", + headers={"Retry-After": str(limiter.retry_after(key))}, + ) + + async def get_current_user(token: str = Depends(oauth2_scheme)) -> dict: """Get current authenticated user from token.""" payload = verify_token(token) @@ -37,8 +58,10 @@ async def get_current_user(token: str = Depends(oauth2_scheme)) -> dict: @router.post("/register", response_model=UserResponse) -async def register(user_data: UserCreate): +async def register(user_data: UserCreate, request: Request = None): """Register a new user.""" + # Throttle mass account creation (each account can trigger billable work). + _rate_limit(register_limiter, f"register:{_client_ip(request)}") async with get_db() as db: # Check if username exists async with db.execute( @@ -77,8 +100,15 @@ async def register(user_data: UserCreate): @router.post("/login", response_model=Token) -async def login(form_data: OAuth2PasswordRequestForm = Depends()): +async def login( + form_data: OAuth2PasswordRequestForm = Depends(), + request: Request = None, +): """Login and get access token.""" + # Throttle online password brute-forcing per (IP, username). + _rate_limit( + login_limiter, f"login:{_client_ip(request)}:{form_data.username.lower()}" + ) async with get_db() as db: async with db.execute( "SELECT id, username, password_hash FROM users WHERE username = ?", diff --git a/backend/app/api/chat.py b/backend/app/api/chat.py index 04ac2be..738e97e 100644 --- a/backend/app/api/chat.py +++ b/backend/app/api/chat.py @@ -2,6 +2,7 @@ import json import logging import math +import time import uuid from typing import List, Dict, Any, AsyncGenerator, Optional from fastapi import APIRouter, Depends, HTTPException @@ -9,6 +10,7 @@ from pydantic import BaseModel, Field from app.api.auth import get_current_user +from app.config import get_settings from app.database import get_db from app.models.chat import ChatRequest, ChatResponse, Conversation from app.services.embedding import get_embedding_service @@ -31,6 +33,37 @@ class FeedbackRequest(BaseModel): note: Optional[str] = Field(default=None, max_length=500) +async def _resolve_conversation_id( + conversation_id: Optional[str], user_id: int, title: str +) -> str: + """Return a conversation_id verified to belong to ``user_id``, creating one if absent. + + SECURITY: when a ``conversation_id`` is supplied we must confirm ownership + before writing. Without this check any authenticated user could inject + messages into another user's conversation and — because the chat path loads + conversation history into the LLM prompt — read that user's content back + through the model's answer. A missing/foreign conversation 404s. + """ + if conversation_id: + async with get_db() as db: + async with db.execute( + "SELECT id FROM conversations WHERE id = ? AND user_id = ?", + (conversation_id, user_id), + ) as cursor: + if not await cursor.fetchone(): + raise HTTPException(status_code=404, detail="Conversation not found") + return conversation_id + + new_id = str(uuid.uuid4()) + async with get_db() as db: + await db.execute( + "INSERT INTO conversations (id, user_id, title) VALUES (?, ?, ?)", + (new_id, user_id, title), + ) + await db.commit() + return new_id + + # How many context chunks to expose to the LLM. We cap aggressively so the # citation markers stay readable — 8 distinct [N] tags in a row is already # noisy. The reranker has already trimmed to the top_k above this. @@ -263,17 +296,26 @@ async def build_rag_context( a COMPARISON instruction asking the LLM to structure the answer to highlight cross-document agreements/disagreements. """ - # Query preprocessing + settings = get_settings() + t_start = time.perf_counter() + + # Query preprocessing. Rewriting costs a full LLM round-trip (~1s) that + # BLOCKS every retrieval step after it, so only pay it for queries long + # enough to plausibly benefit — short keyword-style questions (the demo + # common case) go straight to retrieval. Tunable via QUERY_REWRITE_MIN_LEN + # (0 = always rewrite, i.e. the old behavior). search_query = query - if use_query_rewrite: + if use_query_rewrite and len(query.strip()) >= settings.QUERY_REWRITE_MIN_LEN: query_processor = await get_query_processor() rewritten = await query_processor.rewrite_query(query) if rewritten and len(rewritten) > 0: search_query = rewritten + t_rewrite = time.perf_counter() # Get query embedding embedding_service = await get_embedding_service() query_embedding = await embedding_service.embed_single(search_query) + t_embed = time.perf_counter() # Hybrid search chroma = get_chroma_client() @@ -325,18 +367,23 @@ async def build_rag_context( chunk_ids = [r["chunk_id"] for r in rows] bm25.build_user_index(user_id, chunk_contents, chunk_ids) + # Recall budget per retriever. RRF + the reranker only need enough + # candidates to reliably contain the final top_k; RERANK_RECALL_K=25 + # roughly halves the rerank payload (and its latency) vs the old 50. + recall_k = settings.RERANK_RECALL_K + # Vector search (larger recall for fusion) - vector_results = chroma.search(query_embedding, user_id, top_k=50) + vector_results = chroma.search(query_embedding, user_id, top_k=recall_k) # BM25 search - bm25_results = bm25.search(search_query, user_id, top_k=50) + bm25_results = bm25.search(search_query, user_id, top_k=recall_k) # RRF fusion fused_results = reciprocal_rank_fusion( vector_results, bm25_results, k=60, - top_k=50 + top_k=recall_k ) hybrid_chunks = fused_results else: @@ -354,10 +401,12 @@ async def build_rag_context( chunks = merged else: chunks = hybrid_chunks + t_retrieve = time.perf_counter() # Rerank to get top_k most relevant rerank_service = await get_rerank_service() chunks = await rerank_service.rerank(search_query, chunks, top_k=top_k) + t_rerank = time.perf_counter() # Get context chunks all_chunks = [] @@ -383,6 +432,21 @@ async def build_rag_context( if not any(e["name"] == rel["target"] for e in entities): entities.append({"name": rel["target"], "type": "Related"}) + # Per-stage latency breakdown — lets us compare before/after tuning and + # immediately spot which step dominates time-to-first-token. + t_end = time.perf_counter() + logger.info( + "build_rag_context timing: rewrite=%.3fs embed=%.3fs retrieve=%.3fs " + "rerank=%.3fs enrich=%.3fs total=%.3fs (context_chunks=%d)", + t_rewrite - t_start, + t_embed - t_rewrite, + t_retrieve - t_embed, + t_rerank - t_retrieve, + t_end - t_rerank, + t_end - t_start, + len(all_chunks), + ) + return { "chunks": all_chunks, "entities": entities, @@ -398,17 +462,10 @@ async def chat( """Non-streaming chat with RAG.""" user_id = current_user["id"] - # Get or create conversation - if request.conversation_id: - conversation_id = request.conversation_id - else: - conversation_id = str(uuid.uuid4()) - async with get_db() as db: - await db.execute( - "INSERT INTO conversations (id, user_id, title) VALUES (?, ?, ?)", - (conversation_id, user_id, request.message[:50]) - ) - await db.commit() + # Verify ownership (or create) before writing into the conversation. + conversation_id = await _resolve_conversation_id( + request.conversation_id, user_id, request.message[:50] + ) # Save user message async with get_db() as db: @@ -429,15 +486,21 @@ async def chat( else: context = {"chunks": [], "entities": [], "relations": []} - # Get conversation history + # Get conversation history — the 10 MOST RECENT messages, in chronological + # order. We select newest-first (DESC, id as a tie-breaker for same-second + # timestamps) then reverse, so long conversations keep recent context + # instead of the stale opening turns. conversation_history = [] async with get_db() as db: async with db.execute( - "SELECT role, content FROM messages WHERE conversation_id = ? ORDER BY created_at LIMIT 10", + "SELECT role, content FROM messages " + "WHERE conversation_id = ? ORDER BY created_at DESC, id DESC LIMIT 10", (conversation_id,) ) as cursor: rows = await cursor.fetchall() - conversation_history = [{"role": r["role"], "content": r["content"]} for r in rows] + conversation_history = [ + {"role": r["role"], "content": r["content"]} for r in reversed(rows) + ] # Build a numbered citation context for the prompt if context["chunks"]: @@ -491,46 +554,27 @@ async def chat( num_sources=len(citation["sources"]), ) - # Optional follow-up chips — same semantics as the streaming path. - # Failure is non-fatal; the client just renders zero chips. - followups: List[str] = [] - if request.with_followups: - try: - followups = await llm_service.generate_followups( - request.message, response, n=3, - ) - except Exception as e: - logger.warning("generate_followups failed: %s", e) - followups = [] - return { "message": response, "conversation_id": conversation_id, "related_chunks": context["chunks"][:3], "related_entities": context["entities"][:5], "sources": citation["sources"], - "followups": followups, "citation_coverage": coverage, } async def chat_stream_generator( request: ChatRequest, - user_id: int + user_id: int, + conversation_id: str, ) -> AsyncGenerator[str, None]: - """Generator for streaming chat responses.""" - # Get or create conversation - if request.conversation_id: - conversation_id = request.conversation_id - else: - conversation_id = str(uuid.uuid4()) - async with get_db() as db: - await db.execute( - "INSERT INTO conversations (id, user_id, title) VALUES (?, ?, ?)", - (conversation_id, user_id, request.message[:50]) - ) - await db.commit() + """Generator for streaming chat responses. + ``conversation_id`` is resolved and ownership-verified by the caller + (``chat_stream``) BEFORE this generator starts, so a foreign conversation + yields a clean 404 instead of a half-opened SSE stream. + """ # Save user message async with get_db() as db: await db.execute( @@ -551,6 +595,7 @@ async def chat_stream_generator( context = {"chunks": [], "entities": [], "relations": []} # Build a numbered citation context for the prompt + t_cite_start = time.perf_counter() if context["chunks"]: citation = await _build_citation_context( context["chunks"], user_id, comparison_mode=request.compare_mode, @@ -561,6 +606,7 @@ async def chat_stream_generator( "sources": [], "chunk_id_to_index": {}, } + t_cite_end = time.perf_counter() # Push the sources FIRST so the client can render citation chips while # the text is still streaming. Sources are tied to a query, not a @@ -585,10 +631,47 @@ async def chat_stream_generator( # Stream response llm_service = await get_llm_service() full_response = [] - - async for chunk in llm_service.chat_complete_stream(messages): - full_response.append(chunk) - yield f"data: {chunk}\n\n" + thinking_parts = [] + t_stream_start = time.perf_counter() + t_first_byte: Optional[float] = None + + # chat_complete_stream yields (kind, text) tuples: "thinking" frames + # carry the model's reasoning (hybrid-thinking mode only) and get their + # own SSE event so the client renders them in a separate collapsible + # block; the first-byte metric below stays anchored to the first + # "content" chunk so the two modes are directly comparable. + async for kind, text in llm_service.chat_complete_stream( + messages, enable_thinking=request.enable_thinking + ): + if not text: + continue # defensive: never forward None/empty deltas downstream + if kind == "thinking": + thinking_parts.append(text) + yield f"event: thinking\ndata: {json.dumps({'text': text})}\n\n" + continue + if t_first_byte is None: + t_first_byte = time.perf_counter() + full_response.append(text) + # JSON-encode the chunk: a raw newline in the model output would + # otherwise terminate the SSE `data:` field mid-token and corrupt + # the frame (dropping/garbling text on the client). + yield f"data: {json.dumps({'chunk': text})}\n\n" + + t_stream_end = time.perf_counter() + # Full server-side attribution for one turn, pairing the in-context + # build_rag_context line: first_byte here minus that total there is the + # LLM provider's own prefill/queue latency (the part we don't control). + logger.info( + "chat_stream timing: citation=%.3fs first_byte=%.3fs stream=%.3fs " + "turn=%.3fs (context_chunks=%d, answer_chars=%d, thinking_chars=%d)", + t_cite_end - t_cite_start, + (t_first_byte or t_stream_end) - t_stream_start, + t_stream_end - t_stream_start, + t_stream_end - t_cite_start, + len(context["chunks"]), + sum(len(c) for c in full_response), + sum(len(t) for t in thinking_parts), + ) # Save complete response complete_response = "".join(full_response) @@ -614,21 +697,8 @@ async def chat_stream_generator( num_sources=len(citation["sources"]), ) - # Optional follow-up chips — emit as a separate SSE event AFTER the - # body so the client can render the answer first, then surface the - # chips below it. Failure is non-fatal (generate_followups swallows - # its own errors); we still emit the 'done' event either way. - if request.with_followups: - followups: List[str] = [] - try: - followups = await llm_service.generate_followups( - request.message, complete_response, n=3, - ) - except Exception as e: - logger.warning("generate_followups failed: %s", e) - followups = [] - yield f"event: followups\ndata: {json.dumps({'followups': followups})}\n\n" - + # The answer is fully streamed, saved, and coverage is computed — the + # turn is logically complete. yield f"event: done\ndata: {json.dumps({'conversation_id': conversation_id, 'sources': citation['sources'], 'citation_coverage': coverage})}\n\n" @@ -640,8 +710,14 @@ async def chat_stream( """Streaming chat with RAG.""" user_id = current_user["id"] + # Verify ownership (or create) BEFORE opening the stream, so a foreign + # conversation_id returns a clean 404 instead of a half-opened response. + conversation_id = await _resolve_conversation_id( + request.conversation_id, user_id, request.message[:50] + ) + return StreamingResponse( - chat_stream_generator(request, user_id), + chat_stream_generator(request, user_id, conversation_id), media_type="text/event-stream" ) diff --git a/backend/app/api/documents.py b/backend/app/api/documents.py index 8883afb..ec1e638 100644 --- a/backend/app/api/documents.py +++ b/backend/app/api/documents.py @@ -20,6 +20,7 @@ from app.services.bm25 import get_bm25_service from app.services.entity_extractor import get_entity_extractor from app.services.progress_tracker import get_progress_emitter +from app.services.doc_status import DocStatus, set_document_status logger = logging.getLogger(__name__) @@ -78,6 +79,29 @@ def is_allowed_file(filename: str) -> bool: return get_file_extension(filename) in ALLOWED_EXTENSIONS +def _content_matches_extension(content: bytes, ext: str) -> bool: + """Conservative magic-byte check rejecting grossly mislabelled uploads. + + Only formats with strong, unambiguous signatures are enforced (pdf/docx/ + doc); text formats (txt/md/markdown) and anything uncertain pass through to + the parser. This narrows the attack surface — e.g. a zip bomb handed to the + docx unzip path, or an arbitrary binary renamed to .pdf — without rejecting + legitimate files. + """ + ext = (ext or "").lstrip(".").lower() + head = content[:8] + if ext == "pdf": + return b"%PDF" in content[:1024] + if ext in ("docx", "doc"): + # Modern .docx is a ZIP container; legacy .doc is an OLE2 compound doc. + # Accept either signature for both to stay lenient about misnaming. + return ( + head[:4] in (b"PK\x03\x04", b"PK\x05\x06") + or head[:8] == b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1" + ) + return True + + @router.post("/upload", response_model=DocumentResponse) async def upload_document( background_tasks: BackgroundTasks, @@ -95,12 +119,31 @@ async def upload_document( detail=f"File type not allowed. Allowed: {', '.join(ALLOWED_EXTENSIONS)}" ) - # Check file size - file_content = await file.read() - if len(file_content) > settings.MAX_FILE_SIZE: + # Read in bounded chunks and enforce the size limit AS we read. Reading the + # whole body first (`await file.read()`) let an attacker buffer a multi-GB + # upload into memory before the check ran — a trivial DoS given open + # registration. Now memory is capped at MAX_FILE_SIZE. + _chunk_size = 1024 * 1024 # 1 MiB + _buffers = [] + _total = 0 + while True: + _piece = await file.read(_chunk_size) + if not _piece: + break + _total += len(_piece) + if _total > settings.MAX_FILE_SIZE: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"File too large. Max size: {settings.MAX_FILE_SIZE / 1024 / 1024:.1f}MB" + ) + _buffers.append(_piece) + file_content = b"".join(_buffers) + + # Reject content that clearly doesn't match its declared extension. + if not _content_matches_extension(file_content, get_file_extension(file.filename)): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail=f"File too large. Max size: {settings.MAX_FILE_SIZE / 1024 / 1024:.1f}MB" + detail="File content does not match its extension" ) # Generate document ID @@ -120,11 +163,13 @@ async def upload_document( markdown_content, extracted_title = convert_document_to_markdown(file_path, file_ext[1:]) markdown_content = clean_markdown(markdown_content) except Exception as e: - # Clean up file on error + # Clean up file on error. Log the detail server-side; return a generic + # message so internal parser errors don't leak to the client. + logger.error("Document conversion failed for %s: %s", file_path, e, exc_info=True) os.remove(file_path) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Failed to convert document: {str(e)}" + detail="Failed to convert document" ) if not markdown_content.strip(): @@ -137,19 +182,23 @@ async def upload_document( # Extract title title = extracted_title or extract_title_from_markdown(markdown_content) or file.filename - # Save to database + # Save to database. The document starts in 'pending'; the background + # pipeline advances it through the state machine (services/doc_status.py) + # as each durable checkpoint completes. async with get_db() as db: await db.execute( """INSERT INTO documents - (id, user_id, title, file_path, original_filename, file_type) - VALUES (?, ?, ?, ?, ?, ?)""", - (doc_id, user_id, title, file_path, file.filename, file_ext[1:]) + (id, user_id, title, file_path, original_filename, file_type, status) + VALUES (?, ?, ?, ?, ?, ?, ?)""", + (doc_id, user_id, title, file_path, file.filename, file_ext[1:], + DocStatus.PENDING.value) ) await db.commit() # Get created document with timestamp async with db.execute( - "SELECT id, title, original_filename, file_type, created_at FROM documents WHERE id = ?", + "SELECT id, title, original_filename, file_type, created_at, status " + "FROM documents WHERE id = ?", (doc_id,) ) as cursor: doc = await cursor.fetchone() @@ -182,6 +231,7 @@ async def process_document_background(doc_id: str, user_id: int, markdown: str, neo4j = await get_neo4j_client() await neo4j.create_document_node(doc_id, user_id, title) logger.info("Created document node in Neo4j for %s", doc_id) + await set_document_status(doc_id, DocStatus.DOCUMENT_CREATED) await progress.emit_and_save(doc_id, user_id, "document_created", "Document created", {"stage": "document_created"}) # Chunk the document @@ -192,6 +242,10 @@ async def process_document_background(doc_id: str, user_id: int, markdown: str, if not chunks: logger.warning("No chunks created for doc %s", doc_id) + await set_document_status( + doc_id, DocStatus.FAILED, + error_message="No content could be extracted from the document", + ) await progress.emit_and_save(doc_id, user_id, "error", "No content could be extracted from the document", {"stage": "error"}) return @@ -211,6 +265,10 @@ async def process_document_background(doc_id: str, user_id: int, markdown: str, embeddings = await embedding_service.embed_batch(chunk_contents) except EmbeddingServiceError as embed_err: logger.error("Embedding failed for doc %s: %s", doc_id, embed_err, exc_info=True) + await set_document_status( + doc_id, DocStatus.FAILED, + error_message=f"Embedding failed: {embed_err}", + ) await progress.emit_and_save( doc_id, user_id, "error", f"Embedding failed: {embed_err}", @@ -246,12 +304,14 @@ async def process_document_background(doc_id: str, user_id: int, markdown: str, bm25 = get_bm25_service() bm25.add_to_index(user_id, chunk_contents, chunk_ids) - # Store chunks in SQLite + # Store chunks in SQLite. INSERT OR IGNORE keeps a re-run idempotent: + # chunk_ids are deterministic (derived from doc_id + position), so a + # resumed pipeline re-derives the same ids and must not crash on the PK. async with get_db() as db: for chunk in chunks: hierarchy_path_str = ",".join(chunk.hierarchy.path) if chunk.hierarchy.path else "" await db.execute( - """INSERT INTO chunks + """INSERT OR IGNORE INTO chunks (chunk_id, document_id, user_id, content, hierarchy_path, level, prev_chunk_id, next_chunk_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", @@ -262,22 +322,50 @@ async def process_document_background(doc_id: str, user_id: int, markdown: str, await db.commit() logger.info("Stored %d chunks in SQLite for doc %s", len(chunks), doc_id) - # Create chunk nodes in Neo4j and link them + # Durable checkpoint: vectors (Chroma), chunks (SQLite) and the BM25 + # index are all written now. A crash past this point can resume from + # 'indexed' without redoing embedding. + await set_document_status(doc_id, DocStatus.INDEXED) + + # Create chunk nodes + NEXT links in Neo4j in TWO UNWIND batches + # (one round-trip each) instead of the old per-chunk loop, which + # paid 2N+ serial round-trips (~5–15ms apiece) — for a 100-chunk + # doc that was ~2–3s of pure network latency, now ~20ms total. await progress.emit_and_save(doc_id, user_id, "graph", "Building knowledge graph...", {"stage": "graph", "current": 0, "total": len(chunks)}) - for i, chunk in enumerate(chunks): - await neo4j.create_chunk_node( - chunk.chunk_id, doc_id, user_id, chunk.content, - chunk.hierarchy.path, chunk.position.start_line - ) - await neo4j.create_chunk_links( - chunk.chunk_id, - chunk.position.prev_chunk_id, - chunk.position.next_chunk_id - ) - if i % 5 == 0: # Emit progress every 5 chunks - await progress.emit_and_save(doc_id, user_id, "graph", f"Processing chunk {i+1}/{len(chunks)}", {"stage": "graph", "current": i+1, "total": len(chunks), "percent": int((i+1)/len(chunks)*30) + 50}) + t_graph = time.time() + chunk_payloads = [ + { + "chunk_id": chunk.chunk_id, + "content": chunk.content, + "hierarchy_path": chunk.hierarchy.path, + "position": chunk.position.start_line, + } + for chunk in chunks + ] + created_chunks = await neo4j.create_chunk_nodes_batch(doc_id, user_id, chunk_payloads) + + # prev/next pointers are symmetric (chunker sets both directions), + # so dedupe pairs here — MERGE would dedupe server-side too, but + # sending unique pairs halves the transmitted work. + link_pairs: set = set() + for chunk in chunks: + if chunk.position.prev_chunk_id: + link_pairs.add((chunk.position.prev_chunk_id, chunk.chunk_id)) + if chunk.position.next_chunk_id: + link_pairs.add((chunk.chunk_id, chunk.position.next_chunk_id)) + link_payloads = [{"from_id": a, "to_id": b} for a, b in link_pairs] + linked = await neo4j.create_chunk_links_batch(link_payloads) + logger.info( + "Graph writes for doc %s: %d chunk nodes, %d NEXT links in %.2fs", + doc_id, created_chunks, linked, time.time() - t_graph, + ) await progress.emit_and_save(doc_id, user_id, "graph", "Knowledge graph building complete", {"stage": "graph", "current": len(chunks), "total": len(chunks), "percent": 80}) + # Durable checkpoint: chunk nodes/links are in Neo4j. The (expensive) + # LLM entity extraction that follows has not run yet — a crash here + # resumes from 'graphed' and skips straight to extraction. + await set_document_status(doc_id, DocStatus.GRAPHED) + # Extract entities and relations logger.info("Starting entity extraction with LLM for doc %s", doc_id) total_chunks = len(chunks) @@ -422,6 +510,8 @@ async def process_document_background(doc_id: str, user_id: int, markdown: str, "percent": 95 }) + # Terminal success checkpoint: every store is fully written. + await set_document_status(doc_id, DocStatus.READY) logger.info("Document %s background processing completed", doc_id) # Calculate duration @@ -440,6 +530,13 @@ async def process_document_background(doc_id: str, user_id: int, markdown: str, # Log error but don't fail the upload import traceback logger.error("Background processing error for doc %s: %s", doc_id, e, exc_info=True) + # Mark the document failed so it doesn't linger in an in-progress + # checkpoint forever. Wrapped so a status-write failure can't mask the + # original error or prevent the SSE error event below. + try: + await set_document_status(doc_id, DocStatus.FAILED, error_message=str(e)) + except Exception as status_error: + logger.warning("Failed to mark doc %s as failed: %s", doc_id, status_error) # Emit error event try: progress = get_progress_emitter() @@ -474,7 +571,10 @@ async def list_documents( if tag: # Filter at the SQL layer so we don't pull tags for docs we'd # discard. Inner join keeps only docs that actually have the tag. - sql = """SELECT d.id, d.title, d.original_filename, d.file_type, d.created_at + sql = """SELECT d.id, d.title, d.original_filename, d.file_type, + d.created_at, + COALESCE(d.status, 'pending') AS status, + d.error_message FROM documents d INNER JOIN document_tags t ON t.document_id = d.id AND t.user_id = d.user_id @@ -482,7 +582,8 @@ async def list_documents( ORDER BY d.created_at DESC LIMIT ? OFFSET ?""" params = (user_id, tag, limit, skip) else: - sql = """SELECT id, title, original_filename, file_type, created_at + sql = """SELECT id, title, original_filename, file_type, created_at, + COALESCE(status, 'pending') AS status, error_message FROM documents WHERE user_id = ? ORDER BY created_at DESC LIMIT ? OFFSET ?""" params = (user_id, limit, skip) @@ -532,7 +633,8 @@ async def get_document_detail( async with get_db() as db: async with db.execute( - "SELECT id, title, original_filename, file_type, created_at " + "SELECT id, title, original_filename, file_type, created_at, " + "COALESCE(status, 'pending') AS status, error_message " "FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id), ) as cursor: @@ -582,6 +684,8 @@ async def get_document_detail( "original_filename": doc_row["original_filename"], "file_type": doc_row["file_type"], "created_at": doc_row["created_at"], + "status": doc_row["status"], + "error_message": doc_row["error_message"], "tags": tags, }, "stats": { @@ -869,7 +973,7 @@ async def get_document_chunks( logger.error("Error getting chunks for doc %s: %s", doc_id, e, exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Failed to get chunks: {str(e)}" + detail="Failed to get chunks" ) diff --git a/backend/app/api/progress.py b/backend/app/api/progress.py index 142b38a..65df772 100644 --- a/backend/app/api/progress.py +++ b/backend/app/api/progress.py @@ -78,6 +78,30 @@ async def _authenticate_sse( return await _user_from_token(token) +def _sse_error(detail: str, status_code: int) -> StreamingResponse: + """Return a parseable SSE error event so clients see a typed event.""" + return StreamingResponse( + iter([f"data: {json.dumps({'type': 'error', 'error': detail})}\n\n"]), + media_type="text/event-stream", + status_code=status_code, + ) + + +async def _verify_doc_owner(doc_id: str, user_id: int) -> bool: + """Confirm a document belongs to the user. + + SECURITY: without this, any authenticated user could subscribe to another + user's ``doc_id`` and eavesdrop on their processing events (document + titles, extracted entity names, error text). + """ + async with get_db() as db: + async with db.execute( + "SELECT id FROM documents WHERE id = ? AND user_id = ?", + (doc_id, user_id), + ) as cursor: + return await cursor.fetchone() is not None + + @router.get("/api/progress/{doc_id}") async def stream_progress( doc_id: str, @@ -92,15 +116,15 @@ async def stream_progress( which cannot set custom headers). """ try: - await _authenticate_sse(authorization, token) + current_user = await _authenticate_sse(authorization, token) except HTTPException as exc: # Emit a parseable SSE error event so the client's onmessage sees # a typed event instead of an opaque network failure. - return StreamingResponse( - iter([f"data: {json.dumps({'type': 'error', 'error': exc.detail})}\n\n"]), - media_type="text/event-stream", - status_code=exc.status_code, - ) + return _sse_error(exc.detail, exc.status_code) + + # SECURITY: only the owner may stream a document's progress. + if not await _verify_doc_owner(doc_id, current_user["id"]): + return _sse_error("Document not found", status.HTTP_404_NOT_FOUND) emitter = get_progress_emitter() queue = emitter.subscribe(doc_id) @@ -121,7 +145,8 @@ async def event_generator(): except asyncio.CancelledError: pass finally: - emitter.unsubscribe(doc_id) + # Remove ONLY this subscriber's queue so other watchers survive. + emitter.unsubscribe(doc_id, queue) return StreamingResponse( event_generator(), @@ -146,6 +171,13 @@ async def get_progress_history( except HTTPException as exc: return {"error": exc.detail, "history": []} + # SECURITY: 404 for a document the caller doesn't own, rather than + # silently returning an empty history. + if not await _verify_doc_owner(doc_id, current_user["id"]): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="Document not found" + ) + emitter = get_progress_emitter() history = await emitter.get_history(doc_id, current_user["id"]) return {"history": history} diff --git a/backend/app/auth/rate_limit.py b/backend/app/auth/rate_limit.py new file mode 100644 index 0000000..78a88e1 --- /dev/null +++ b/backend/app/auth/rate_limit.py @@ -0,0 +1,57 @@ +"""Lightweight in-memory sliding-window rate limiter (no external dependencies). + +Used to throttle authentication endpoints (login/register) against brute-force +and account-spraying attacks. State lives in process memory, so limits are +per-worker and reset on restart — adequate as a first line of defence for a +single-instance deployment. Swap for a shared store (e.g. Redis) if the app +ever runs multiple workers behind a load balancer. +""" +import time +from collections import defaultdict, deque +from typing import Deque, Dict + + +class SlidingWindowLimiter: + """Allow at most ``max_calls`` hits per ``key`` within ``window_seconds``.""" + + def __init__(self, max_calls: int, window_seconds: int) -> None: + self.max_calls = max_calls + self.window_seconds = window_seconds + self._hits: Dict[str, Deque[float]] = defaultdict(deque) + + def _purge(self, key: str, now: float) -> Deque[float]: + dq = self._hits[key] + cutoff = now - self.window_seconds + while dq and dq[0] <= cutoff: + dq.popleft() + return dq + + def is_allowed(self, key: str) -> bool: + """Record a hit for ``key`` and return whether it is within the limit.""" + now = time.monotonic() + dq = self._purge(key, now) + if len(dq) >= self.max_calls: + return False + dq.append(now) + # Opportunistic memory guard: drop fully-expired keys so an attacker + # spraying unique identifiers can't grow the dict without bound. + if len(self._hits) > 10000: + for k in [k for k, v in self._hits.items() if not v]: + del self._hits[k] + return True + + def retry_after(self, key: str) -> int: + """Seconds until the oldest hit for ``key`` leaves the window.""" + dq = self._hits.get(key) + if not dq: + return 0 + elapsed = time.monotonic() - dq[0] + return max(1, int(self.window_seconds - elapsed) + 1) + + +# Login: a handful of attempts per (IP, username) per minute stops online +# brute force while tolerating honest mistypes. +login_limiter = SlidingWindowLimiter(max_calls=8, window_seconds=60) +# Registration: tight per-IP to blunt mass account creation (each account +# triggers billable LLM/embedding work on first upload). +register_limiter = SlidingWindowLimiter(max_calls=10, window_seconds=3600) diff --git a/backend/app/config.py b/backend/app/config.py index 7e2928e..f616dee 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -28,7 +28,7 @@ class Settings(BaseSettings): # Bailian (百炼) LLM - 使用 OpenAI 兼容模式 BAILIAN_API_KEY: str = "" BAILIAN_BASE_URL: str = "https://dashscope.aliyuncs.com/compatible-mode/v1" - BAILIAN_MODEL: str = "qwen-flash" + BAILIAN_MODEL: str = "qwen3.7-flash" # LLM Settings LLM_MODEL_KIMI: str = "kimi-k2-0905-preview" @@ -50,11 +50,34 @@ class Settings(BaseSettings): # Rerank RERANK_MODEL: str = "Qwen/Qwen3-Reranker-8B" + # Retrieval latency tuning + # Query rewriting costs a full LLM round-trip that BLOCKS retrieval, so + # it is only worth paying for longer queries. Queries shorter than this + # (stripped char count) skip the rewrite entirely. Set to 0 to always + # rewrite (old behavior), or to a very large number to never rewrite. + QUERY_REWRITE_MIN_LEN: int = 20 + # Candidates each retriever (vector + BM25) feeds into RRF and then the + # reranker. The reranker only needs enough candidates to reliably contain + # the final top_k; 25 roughly halves rerank payload/latency vs 50 with no + # measurable hit to top-5 quality. + RERANK_RECALL_K: int = 25 + # Entity Extraction ENABLE_LLM_EXTRACTION: bool = True USE_RULE_EXTRACTION: bool = False # 纯 LLM 模式，不使用规则提取（更快） ENTITY_BATCH_SIZE: int = 200 # 实体提取批次大小 ENTITY_EXTRACTION_DELAY: float = 0 + # Concurrent in-flight LLM extraction requests. The old hard-coded 50 + # regularly tripped provider rate limits (429s), and each 429 paid a + # 2–4s backoff — so a lower ceiling sustains HIGHER effective + # throughput. 20 keeps ~20 chunks extracting in parallel without + # provoking 429 storms on the Bailian compatible-mode tier. + LLM_EXTRACTION_CONCURRENCY: int = 20 + # max_tokens for extraction calls. Entity/relation JSON for a 2000-char + # chunk fits comfortably under 1024 tokens; the old 8000 default made + # the provider reserve/allocate far more generation budget than any + # extraction response could ever use. + LLM_EXTRACT_MAX_TOKENS: int = 1024 # CORS - comma-separated list of allowed origins (no wildcards with credentials) CORS_ALLOWED_ORIGINS: str = "http://localhost:5173,http://localhost:3000,http://127.0.0.1:5173" @@ -68,19 +91,31 @@ class Config: extra = "ignore" -_DEFAULT_JWT_SECRET = "your-secret-key-change-this-in-production" +# Publicly-known placeholder secrets that must never be used at runtime. Both +# the code default and the .env.example placeholder are public, so accepting +# either would let anyone forge valid JWTs for any user (full auth bypass). +_INSECURE_JWT_SECRETS = { + "", + "your-secret-key-change-this-in-production", + "replace-me-with-a-strong-random-value", +} @lru_cache() def get_settings() -> Settings: - """Get cached settings instance. Refuses to start with default JWT_SECRET in production.""" + """Get cached settings instance. + + Refuses to start with a known placeholder JWT_SECRET in ANY environment — + the placeholders are public and would allow trivial authentication bypass. + Generate one with: + python -c "import secrets; print(secrets.token_urlsafe(48))" + """ settings = Settings() - if settings.JWT_SECRET == _DEFAULT_JWT_SECRET: - import os - env = os.environ.get("APP_ENV", "development").lower() - if env in ("production", "prod"): - raise RuntimeError( - "JWT_SECRET is set to the default placeholder in production. " - "Set a strong random value via the JWT_SECRET environment variable." - ) + if settings.JWT_SECRET in _INSECURE_JWT_SECRETS: + raise RuntimeError( + "JWT_SECRET is set to an insecure placeholder. Generate a strong " + "random value via `python -c \"import secrets; " + "print(secrets.token_urlsafe(48))\"` and set it with the JWT_SECRET " + "environment variable." + ) return settings diff --git a/backend/app/database.py b/backend/app/database.py index 54b1ddf..f870bb0 100644 --- a/backend/app/database.py +++ b/backend/app/database.py @@ -7,6 +7,32 @@ from app.config import get_settings +async def _ensure_document_status_columns(db) -> None: + """Idempotently add the state-machine columns to an existing documents table. + + ``CREATE TABLE IF NOT EXISTS`` never alters a table that already exists, so + databases created before the processing state machine was introduced lack + these columns. We add them and backfill legacy rows as ``ready`` — those + documents were processed under the old system and already live in the + stores, so they must not surface as stuck-in-``pending``. + """ + async with db.execute("PRAGMA table_info(documents)") as cursor: + existing = {row[1] for row in await cursor.fetchall()} + + if "status" not in existing: + await db.execute("ALTER TABLE documents ADD COLUMN status TEXT") + await db.execute( + "UPDATE documents SET status = 'ready' WHERE status IS NULL" + ) + if "error_message" not in existing: + await db.execute("ALTER TABLE documents ADD COLUMN error_message TEXT") + if "updated_at" not in existing: + await db.execute("ALTER TABLE documents ADD COLUMN updated_at TIMESTAMP") + await db.execute( + "UPDATE documents SET updated_at = created_at WHERE updated_at IS NULL" + ) + + async def init_db(): """Initialize SQLite database with required tables.""" settings = get_settings() @@ -15,6 +41,9 @@ async def init_db(): os.makedirs(os.path.dirname(settings.SQLITE_PATH), exist_ok=True) async with aiosqlite.connect(settings.SQLITE_PATH) as db: + # Enforce foreign keys (SQLite keeps them OFF by default) so the + # ON DELETE CASCADE rules declared below actually take effect. + await db.execute("PRAGMA foreign_keys = ON") # Create users table await db.execute(""" CREATE TABLE IF NOT EXISTS users ( @@ -25,7 +54,12 @@ async def init_db(): ) """) - # Create documents table + # Create documents table. + # `status` is the processing state machine (see services/doc_status.py): + # pending -> document_created -> indexed -> graphed -> ready, with + # `failed` reachable from any non-terminal state. Defaults to 'pending' + # so a row that never gets processed is visibly unprocessed rather than + # silently treated as done. await db.execute(""" CREATE TABLE IF NOT EXISTS documents ( id TEXT PRIMARY KEY, @@ -34,10 +68,15 @@ async def init_db(): file_path TEXT, original_filename TEXT, file_type TEXT, + status TEXT NOT NULL DEFAULT 'pending', + error_message TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE ) """) + # Migrate databases created before the state machine existed. + await _ensure_document_status_columns(db) # Create chunks table for tracking await db.execute(""" @@ -162,4 +201,7 @@ async def get_db() -> AsyncGenerator[aiosqlite.Connection, None]: settings = get_settings() async with aiosqlite.connect(settings.SQLITE_PATH) as db: db.row_factory = aiosqlite.Row + # PRAGMA is per-connection; enable FK enforcement on every connection + # so cascade deletes (messages, tags, feedback) work app-wide. + await db.execute("PRAGMA foreign_keys = ON") yield db diff --git a/backend/app/main.py b/backend/app/main.py index 263fb7b..601fb2d 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -43,6 +43,14 @@ async def lifespan(app: FastAPI): chroma = get_chroma_client() chroma.close() + # Close lazily-initialized LLM/rerank HTTP clients (no-op if never used). + from app.services.llm import close_llm_service + from app.services.reranker import close_rerank_service + from app.services.embedding import close_embedding_service + await close_llm_service() + await close_rerank_service() + await close_embedding_service() + logger.info("Knowledge Graph System Stopped") @@ -79,8 +87,6 @@ def create_app() -> FastAPI: app.include_router(tags.router) app.include_router(timeline.router) app.include_router(dashboard.router) - app.include_router(timeline.router) - app.include_router(tags.router) @app.get("/") async def root(): diff --git a/backend/app/models/chat.py b/backend/app/models/chat.py index c4bd8ea..0523341 100644 --- a/backend/app/models/chat.py +++ b/backend/app/models/chat.py @@ -17,7 +17,11 @@ class ChatRequest(BaseModel): include_context: bool = True use_graph_rag: bool = False # opt-in: graph-first candidate set compare_mode: bool = False # opt-in: structure answer as a comparison - with_followups: bool = True # opt-out: skip generating follow-up chips + # opt-in: Qwen hybrid-thinking mode — the model streams its reasoning + # (forwarded as `event: thinking` SSE frames) before the answer body. + # Noticeably slower first token, occasionally better answers. Only the + # streaming endpoint honors it; the non-streaming /chat path ignores it. + enable_thinking: bool = False class ChatResponse(BaseModel): @@ -26,10 +30,6 @@ class ChatResponse(BaseModel): conversation_id: str related_chunks: List[Dict[str, Any]] = [] related_entities: List[Dict[str, Any]] = [] - # Up to 3 follow-up question chips generated by the LLM after the - # answer. Empty list = no followups (either opted out or the LLM - # call failed — both are non-fatal). - followups: List[str] = [] # Fraction (0.0–1.0) of source chips that the answer actually # cites. A low value (e.g. 0.2) means the LLM mostly hand-waved # and the sources weren't really used. Front-end renders this as diff --git a/backend/app/models/document.py b/backend/app/models/document.py index 7302aef..01c1b77 100644 --- a/backend/app/models/document.py +++ b/backend/app/models/document.py @@ -22,6 +22,11 @@ class DocumentInDB(DocumentBase): original_filename: str file_type: str created_at: datetime + # Processing state machine (services/doc_status.py). Optional so rows + # predating the columns (or partial reads) still validate. + status: Optional[str] = "pending" + error_message: Optional[str] = None + updated_at: Optional[datetime] = None class Config: from_attributes = True @@ -33,6 +38,12 @@ class DocumentResponse(DocumentBase): original_filename: str file_type: str created_at: datetime + # Processing state machine checkpoint (services/doc_status.py): one of + # pending / document_created / indexed / graphed / ready / failed. Lets the + # UI show a stuck or failed document without polling the SSE history. + # Defaults to 'pending' for payloads that predate the field. + status: str = "pending" + error_message: Optional[str] = None # Tags are surfaced on every document-list response so the UI doesn't # have to round-trip per card. Always present (possibly empty) — never # null — so the client can iterate without null-checks. diff --git a/backend/app/models/user.py b/backend/app/models/user.py index 35e4124..11d778e 100644 --- a/backend/app/models/user.py +++ b/backend/app/models/user.py @@ -7,6 +7,10 @@ _USERNAME_RE = re.compile(r"^[A-Za-z0-9_.-]+$") +# bcrypt only hashes the first 72 BYTES of a password; anything beyond is +# silently ignored. Enforce the byte limit at the API layer so two passwords +# differing only past byte 72 can't collide to the same hash. +_BCRYPT_MAX_BYTES = 72 class UserBase(BaseModel): @@ -25,7 +29,7 @@ def _validate_username(cls, value: str) -> str: class UserCreate(UserBase): """User creation model.""" - password: str = Field(..., min_length=8, max_length=100) + password: str = Field(..., min_length=8, max_length=72) @field_validator("password") @classmethod @@ -34,6 +38,10 @@ def _validate_password(cls, value: str) -> str: raise ValueError("password must contain at least one letter") if not re.search(r"\d", value): raise ValueError("password must contain at least one digit") + if len(value.encode("utf-8")) > _BCRYPT_MAX_BYTES: + raise ValueError( + "password is too long (max 72 bytes after UTF-8 encoding)" + ) return value diff --git a/backend/app/services/chroma_client.py b/backend/app/services/chroma_client.py index 11fb4f4..0ade815 100644 --- a/backend/app/services/chroma_client.py +++ b/backend/app/services/chroma_client.py @@ -54,7 +54,10 @@ def add_chunks( cleaned[key] = value cleaned_metadatas.append(cleaned) - self._collection.add( + # upsert (not add) so a resumed/re-run ingestion is idempotent: chunk + # ids are deterministic, and `add` on an existing id only warns and + # skips, whereas upsert guarantees the stored vector matches the run. + self._collection.upsert( ids=chunk_ids, documents=documents, embeddings=embeddings, diff --git a/backend/app/services/doc_status.py b/backend/app/services/doc_status.py new file mode 100644 index 0000000..a0ae96f --- /dev/null +++ b/backend/app/services/doc_status.py @@ -0,0 +1,202 @@ +"""Document processing state machine. + +The ``documents.status`` column is the single source of truth for "how far has +this document been processed?". Unlike ``progress_history`` (a fine-grained +event log for the progress UI), ``status`` records only *durable checkpoints* — +states whose work has actually been persisted to the backing stores (Neo4j / +Chroma / SQLite). That distinction is deliberate: it is what makes future +crash-recovery / resume possible, because on restart we can read the checkpoint +and skip work that is already durable. + +Lifecycle (monotonic, forward-only):: + + pending -> document_created -> indexed -> graphed -> ready + (any non-terminal state) -> failed + +Transition rules enforced by :func:`set_document_status`: + +* forward move -> applied +* same state -> idempotent no-op (returns False) +* backward move -> idempotent no-op (resume/retry safe, returns False) +* non-terminal -> failed-> applied, records ``error_message`` +* leaving a terminal -> raises :class:`InvalidStatusTransition` + (``ready``/``failed``) (use :func:`reset_for_retry` to reprocess a failed doc) +""" +import logging +from enum import Enum +from typing import Optional + +from app.database import get_db + +logger = logging.getLogger(__name__) + + +class DocStatus(str, Enum): + """Durable checkpoints in a document's processing lifecycle.""" + + PENDING = "pending" # row created; background work not started + DOCUMENT_CREATED = "document_created" # Neo4j Document node exists + INDEXED = "indexed" # chunks in Chroma + SQLite (+ BM25) + GRAPHED = "graphed" # chunk nodes/links in Neo4j + READY = "ready" # entities + relations saved; complete + FAILED = "failed" # terminal error (see error_message) + + +# Forward ordering of the non-terminal checkpoint states. ``failed`` is handled +# separately because it is reachable from anywhere and has no rank. +_RANK = { + DocStatus.PENDING: 0, + DocStatus.DOCUMENT_CREATED: 1, + DocStatus.INDEXED: 2, + DocStatus.GRAPHED: 3, + DocStatus.READY: 4, +} + +#: States from which no further transition is allowed without an explicit reset. +TERMINAL_STATES = frozenset({DocStatus.READY, DocStatus.FAILED}) + + +class InvalidStatusTransition(Exception): + """Raised when a transition would leave a terminal state.""" + + def __init__(self, current: str, target: str) -> None: + self.current = current + self.target = target + super().__init__( + f"illegal document status transition: {current!r} -> {target!r}" + ) + + +class DocumentNotFound(Exception): + """Raised when a status update targets a non-existent document.""" + + def __init__(self, doc_id: str) -> None: + self.doc_id = doc_id + super().__init__(f"document not found: {doc_id!r}") + + +def _coerce(status) -> DocStatus: + """Accept a DocStatus or its string value; reject unknowns loudly.""" + return DocStatus(status) + + +def validate_transition(current, target) -> bool: + """Return True if ``current -> target`` should be written. + + Raises :class:`InvalidStatusTransition` for moves out of a terminal state. + Returns False for idempotent no-ops (same or backward move) so the caller + can skip the write. + """ + current = _coerce(current) + target = _coerce(target) + + if current == target: + return False # idempotent + + if current is DocStatus.FAILED: + # Re-failing is a no-op; anything else needs an explicit retry reset. + if target is DocStatus.FAILED: + return False + raise InvalidStatusTransition(current.value, target.value) + + if current is DocStatus.READY: + raise InvalidStatusTransition(current.value, target.value) + + if target is DocStatus.FAILED: + return True # any non-terminal state may fail + + # Both are ranked checkpoint states: forward applies, backward is a no-op. + if _RANK[target] <= _RANK[current]: + logger.debug( + "ignoring backward document status move %s -> %s", + current.value, target.value, + ) + return False + return True + + +async def set_document_status( + doc_id: str, + status, + *, + error_message: Optional[str] = None, +) -> bool: + """Idempotently move a document to ``status``. + + Returns True if the row was updated, False if it was a no-op. Clears + ``error_message`` on any successful (non-failed) transition; records it on + a transition to ``failed``. + """ + target = _coerce(status) + + async with get_db() as db: + async with db.execute( + "SELECT status FROM documents WHERE id = ?", (doc_id,) + ) as cursor: + row = await cursor.fetchone() + if row is None: + raise DocumentNotFound(doc_id) + + current = row["status"] or DocStatus.PENDING.value + if not validate_transition(current, target): + return False + + if target is DocStatus.FAILED: + await db.execute( + "UPDATE documents " + "SET status = ?, error_message = ?, updated_at = CURRENT_TIMESTAMP " + "WHERE id = ?", + (target.value, error_message, doc_id), + ) + else: + await db.execute( + "UPDATE documents " + "SET status = ?, error_message = NULL, updated_at = CURRENT_TIMESTAMP " + "WHERE id = ?", + (target.value, doc_id), + ) + await db.commit() + + logger.info("document %s status -> %s", doc_id, target.value) + return True + + +async def reset_for_retry(doc_id: str) -> bool: + """Move a ``failed`` document back to ``pending`` so it can be reprocessed. + + Only valid from ``failed``; raises :class:`InvalidStatusTransition` + otherwise (a ``ready`` document is not retried via this path). + """ + async with get_db() as db: + async with db.execute( + "SELECT status FROM documents WHERE id = ?", (doc_id,) + ) as cursor: + row = await cursor.fetchone() + if row is None: + raise DocumentNotFound(doc_id) + + current = _coerce(row["status"] or DocStatus.PENDING.value) + if current is not DocStatus.FAILED: + raise InvalidStatusTransition(current.value, DocStatus.PENDING.value) + + await db.execute( + "UPDATE documents " + "SET status = ?, error_message = NULL, updated_at = CURRENT_TIMESTAMP " + "WHERE id = ?", + (DocStatus.PENDING.value, doc_id), + ) + await db.commit() + logger.info("document %s reset for retry", doc_id) + return True + + +async def get_document_status(doc_id: str) -> Optional[str]: + """Return the current status string for a document, or None if missing.""" + async with get_db() as db: + async with db.execute( + "SELECT status FROM documents WHERE id = ?", (doc_id,) + ) as cursor: + row = await cursor.fetchone() + if row is None: + return None + return row["status"] or DocStatus.PENDING.value diff --git a/backend/app/services/embedding.py b/backend/app/services/embedding.py index 3cbbe14..570af22 100644 --- a/backend/app/services/embedding.py +++ b/backend/app/services/embedding.py @@ -67,7 +67,9 @@ class EmbeddingService: Reliability features: - Exponential backoff on transport / 5xx errors (5 attempts: 1,2,4,8,16s). - - Fresh httpx.AsyncClient per attempt (avoids keep-alive stale-connection reuse). + - Shared keep-alive httpx.AsyncClient (one TLS handshake, reused across + calls; a stale connection surfaces as RemoteProtocolError, which the + retry policy handles). - 4xx errors are NOT retried — they are surfaced immediately. - Self-healing cache: rows whose bytes do not look like JSON are deleted on read, so legacy pickle data cannot poison the cache forever. @@ -81,6 +83,29 @@ def __init__(self): self.api_key = self.settings.SILICON_FLOW_API_KEY self.model = self.settings.EMBEDDING_MODEL self._semaphore = asyncio.Semaphore(5) + self._client: Optional[httpx.AsyncClient] = None + + async def _get_client(self) -> httpx.AsyncClient: + """Get or create the shared HTTP client. + + Reusing one client keeps the TLS connection to SiliconFlow warm + across calls — the old per-attempt client paid a fresh TCP+TLS + handshake (~100–300ms) on EVERY embedding request, which dominated + latency for single-text query embeddings and serialized batch + uploads alike. + """ + if self._client is None: + self._client = httpx.AsyncClient( + timeout=REQUEST_TIMEOUT_SECONDS, + limits=httpx.Limits(max_connections=20, max_keepalive_connections=10), + ) + return self._client + + async def close(self) -> None: + """Close the shared HTTP client (no-op if never created).""" + if self._client is not None: + await self._client.aclose() + self._client = None async def _delete_corrupt_cache_row(self, db, text_hash: str, reason: str) -> None: """Remove a single corrupt cache row.""" @@ -144,51 +169,51 @@ async def _call_with_retry(self, payload: dict) -> dict: last_error: Optional[BaseException] = None url = f"{self.base_url}/embeddings" headers = {"Authorization": f"Bearer {self.api_key}"} + client = await self._get_client() for attempt in range(MAX_ATTEMPTS): - async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT_SECONDS) as client: - try: - response = await client.post(url, headers=headers, json=payload) - except RETRYABLE_EXCEPTIONS as e: - last_error = e - if attempt < MAX_ATTEMPTS - 1: - delay = RETRY_DELAYS_SECONDS[attempt] - logger.warning( - "Embedding call %d/%d transport error: %s — retrying in %ds", - attempt + 1, MAX_ATTEMPTS, e, delay, - ) - await asyncio.sleep(delay) - continue - logger.error("Embedding call failed after %d attempts: %s", MAX_ATTEMPTS, e) - raise EmbeddingServiceError( - f"SiliconFlow unreachable after {MAX_ATTEMPTS} attempts: {e}" - ) from e - - if response.status_code in RETRYABLE_STATUS_CODES: - last_error = httpx.HTTPStatusError( - f"status {response.status_code}", - request=response.request, - response=response, + try: + response = await client.post(url, headers=headers, json=payload) + except RETRYABLE_EXCEPTIONS as e: + last_error = e + if attempt < MAX_ATTEMPTS - 1: + delay = RETRY_DELAYS_SECONDS[attempt] + logger.warning( + "Embedding call %d/%d transport error: %s — retrying in %ds", + attempt + 1, MAX_ATTEMPTS, e, delay, ) - if attempt < MAX_ATTEMPTS - 1: - delay = RETRY_DELAYS_SECONDS[attempt] - logger.warning( - "Embedding call %d/%d got HTTP %d — retrying in %ds", - attempt + 1, MAX_ATTEMPTS, response.status_code, delay, - ) - await asyncio.sleep(delay) - continue - raise EmbeddingServiceError( - f"SiliconFlow returned {response.status_code} after {MAX_ATTEMPTS} attempts" - ) from last_error - - if response.status_code >= 400: - body_preview = response.text[:300] if response.text else "" - raise EmbeddingServiceError( - f"SiliconFlow rejected request (HTTP {response.status_code}): {body_preview}" + await asyncio.sleep(delay) + continue + logger.error("Embedding call failed after %d attempts: %s", MAX_ATTEMPTS, e) + raise EmbeddingServiceError( + f"SiliconFlow unreachable after {MAX_ATTEMPTS} attempts: {e}" + ) from e + + if response.status_code in RETRYABLE_STATUS_CODES: + last_error = httpx.HTTPStatusError( + f"status {response.status_code}", + request=response.request, + response=response, + ) + if attempt < MAX_ATTEMPTS - 1: + delay = RETRY_DELAYS_SECONDS[attempt] + logger.warning( + "Embedding call %d/%d got HTTP %d — retrying in %ds", + attempt + 1, MAX_ATTEMPTS, response.status_code, delay, ) + await asyncio.sleep(delay) + continue + raise EmbeddingServiceError( + f"SiliconFlow returned {response.status_code} after {MAX_ATTEMPTS} attempts" + ) from last_error - return response.json() + if response.status_code >= 400: + body_preview = response.text[:300] if response.text else "" + raise EmbeddingServiceError( + f"SiliconFlow rejected request (HTTP {response.status_code}): {body_preview}" + ) + + return response.json() raise EmbeddingServiceError( f"Embedding call failed after {MAX_ATTEMPTS} attempts: {last_error}" @@ -305,3 +330,11 @@ async def get_embedding_service() -> EmbeddingService: if _embedding_service is None: _embedding_service = EmbeddingService() return _embedding_service + + +async def close_embedding_service() -> None: + """Close the shared embedding HTTP client at shutdown (no-op if never created).""" + global _embedding_service + if _embedding_service is not None: + await _embedding_service.close() + _embedding_service = None diff --git a/backend/app/services/entity_extractor.py b/backend/app/services/entity_extractor.py index b7be6f3..808b1a4 100644 --- a/backend/app/services/entity_extractor.py +++ b/backend/app/services/entity_extractor.py @@ -1,5 +1,4 @@ """Entity and relation extraction service combining rules and LLM.""" -import asyncio import logging import re from dataclasses import dataclass @@ -218,52 +217,71 @@ async def extract_relations(self, text: str, entities: List[ExtractedEntity]) -> print(f" Total relations extracted: {len(relations)}") return relations - async def _extract_entities_batch_optimized( + async def _extract_entities_and_relations_llm( self, chunks: List[Any] - ) -> Dict[str, List[ExtractedEntity]]: - """Extract entities from chunks using LLM with batch processing.""" - chunk_entities = {} - batch_size = self.settings.ENTITY_BATCH_SIZE # 使用配置的批次大小 + ) -> Dict[str, Dict[str, List]]: + """Extract entities AND relations for every chunk via ONE LLM call + per chunk — see LLMService.extract_entities_and_relations_batch for + why this replaces the old two-stage design (2N calls + a barrier + where no relation call could start until every chunk's entities + were back). + + Returns {chunk_id: {"entities": [ExtractedEntity], + "relations": [ExtractedRelation]}}. + + On total failure every chunk maps to empty lists, so the pipeline + degrades to "document indexed without a graph" instead of failing + the upload. + """ llm_service = await get_llm_service() + texts = [c.content for c in chunks] + try: + raw_results = await llm_service.extract_entities_and_relations_batch(texts) + except Exception as e: + logger.warning( + "LLM combined extraction failed; returning empty results: %s", + e, exc_info=True, + ) + return {c.chunk_id: {"entities": [], "relations": []} for c in chunks} + + results: Dict[str, Dict[str, List]] = {} + for chunk, raw in zip(chunks, raw_results): + entity_dicts = raw.get("entities", []) if isinstance(raw, dict) else [] + relation_dicts = raw.get("relations", []) if isinstance(raw, dict) else [] + + entities = [ + ExtractedEntity( + name=str(e.get("name", "")).strip(), + type=e.get("type") or "OTHER", + description=e.get("description"), + source="llm", + ) + for e in entity_dicts + if str(e.get("name") or "").strip() + ] + + # Relations may only reference entities extracted from THIS + # chunk. Anything else would be silently dropped later by + # create_relations_batch's MATCH-by-name — filtering here keeps + # the logged counts honest. + known_names = {e.name for e in entities} + relations: List[ExtractedRelation] = [] + for r in relation_dicts: + source = str(r.get("source") or "").strip() + target = str(r.get("target") or "").strip() + if not source or not target or source == target: + continue + if source not in known_names or target not in known_names: + continue + relations.append(ExtractedRelation( + source=source, + target=target, + relation_type=r.get("relation_type") or "MENTIONS", + relation_source="llm", + )) - # Split into batches - batches = [] - for i in range(0, len(chunks), batch_size): - batches.append(chunks[i:i + batch_size]) - - print(f" Processing {len(chunks)} chunks in {len(batches)} batches (batch_size={batch_size})") - - async def _process_batch(batch: List[Any]) -> Dict[str, List[ExtractedEntity]]: - result = {} - texts = [c.content for c in batch] - try: - # 批量调用 LLM，一次处理多文本 - llm_results = await llm_service.extract_entities_batch(texts) - for chunk, entities_data in zip(batch, llm_results): - result[chunk.chunk_id] = [ - ExtractedEntity( - name=e.get("name", ""), - type=e.get("type", "OTHER"), - description=e.get("description"), - source="llm" - ) - for e in entities_data if e.get("name") - ] - except Exception as e: - logger.warning("LLM entity batch failed; returning empty results: %s", e, exc_info=True) - for chunk in batch: - result[chunk.chunk_id] = [] - return result - - # Run all batches concurrently - tasks = [_process_batch(batch) for batch in batches] - results = await asyncio.gather(*tasks, return_exceptions=True) - - for result in results: - if isinstance(result, dict): - chunk_entities.update(result) - - return chunk_entities + results[chunk.chunk_id] = {"entities": entities, "relations": relations} + return results def _merge_entity_results( self, rule_entities: List[ExtractedEntity], llm_entities: List[ExtractedEntity] @@ -281,170 +299,47 @@ def _merge_entity_results( return list(entity_dict.values()) - async def _extract_cooccurrence_relations( - self, text: str, entities: List[ExtractedEntity] - ) -> List[ExtractedRelation]: - """Extract relations based on co-occurrence distance.""" - relations = [] - if len(entities) < 2: - return relations - - entity_names = [e.name for e in entities] - entity_positions = {name: [] for name in entity_names} - - for name in entity_names: - try: - for match in re.finditer(re.escape(name), text): - entity_positions[name].append(match.start()) - except re.error: - continue - - created_pairs = set() - for name1 in entity_names: - for name2 in entity_names: - if name1 >= name2: - continue - - pair_key = tuple(sorted([name1, name2])) - if pair_key in created_pairs: - continue - - positions1 = entity_positions.get(name1, []) - positions2 = entity_positions.get(name2, []) - - if not positions1 or not positions2: - continue - - for pos1 in positions1: - for pos2 in positions2: - if abs(pos1 - pos2) < 300: - relations.append(ExtractedRelation( - source=name1, - target=name2, - relation_type="MENTIONS" - )) - created_pairs.add(pair_key) - break - else: - continue - break - - return relations - - async def _extract_relations_llm_parallel( - self, - chunks: List[Any], - chunk_entities: List[Dict], - llm_service - ) -> List[ExtractedRelation]: - """Extract relations using LLM with batch processing.""" - print(f" Extracting relations with LLM (batch mode)...") - - all_relations = [] - batch_size = self.settings.ENTITY_BATCH_SIZE - - # Split into batches - rel_batches = [] - for i in range(0, len(chunks), batch_size): - batch_chunks = chunks[i:i + batch_size] - batch_chunk_entities = chunk_entities[i:i + batch_size] - rel_batches.append((batch_chunks, batch_chunk_entities)) - - print(f" Processing {len(chunks)} chunks in {len(rel_batches)} relation batches") - - async def _extract_relations_batch(batch_chunks, batch_ce): - texts = [c.content for c in batch_chunks] - ents_list = [[{"name": e.name, "type": e.type} for e in cd["entities"]] - for cd in batch_ce] - return await llm_service.extract_relations_batch(texts, ents_list) - - # Run all batches concurrently - tasks = [_extract_relations_batch(bc, bce) for bc, bce in rel_batches] - results = await asyncio.gather(*tasks, return_exceptions=True) - - for result in results: - if isinstance(result, Exception): - print(f" Relation batch error: {result}") - continue - for relations_data in result: - for rel in relations_data: - all_relations.append(ExtractedRelation( - source=rel.get("source", ""), - target=rel.get("target", ""), - relation_type=rel.get("relation_type", "MENTIONS"), - relation_source="llm" - )) - - print(f" Total relations extracted: {len(all_relations)}") - return all_relations - - async def _extract_relations_cooccurrence_parallel( - self, - chunks: List[Any], - chunk_entities: List[Dict] - ) -> List[ExtractedRelation]: - """并行提取关系 - 共现模式""" - print(f" Extracting relations with co-occurrence (parallel)...") - - all_relations = [] - - # 并行处理所有 chunk - tasks = [ - self._extract_cooccurrence_relations(chunk.content, chunk_data["entities"]) - for chunk, chunk_data in zip(chunks, chunk_entities) - ] - - results = await asyncio.gather(*tasks, return_exceptions=True) - - for result in results: - if isinstance(result, Exception): - continue - all_relations.extend(result) - - print(f" Total relations extracted: {len(all_relations)}") - return all_relations - async def process_chunks(self, chunks: List[Any], use_rule_extraction: bool = False) -> Dict[str, Any]: """Process multiple chunks to extract entities and relations. Args: chunks: List of chunks to process - use_rule_extraction: If True, use rule-based extraction first then LLM. + use_rule_extraction: If True, merge rule-based entities in first. If False, use LLM only (faster, recommended). 流程： 1. 规则提取（可选）- 快速获得基础实体 - 2. LLM 实体提取 - 并发处理所有 chunks - 3. LLM 关系提取 - 基于提取的实体 + 2. 每 chunk 一次合并 LLM 调用，同时返回实体和关系 + （旧设计是「N 次实体调用 → stage barrier → N 次关系调用」， + 调用数翻倍且关系抽取必须等全部实体返回；合并后 LLM 往返 + 减半、全程并发重叠） + 3. 合并规则 + LLM 实体并去重 """ - print(f" Processing {len(chunks)} chunks for entity extraction...") + logger.info("Processing %d chunks for entity extraction...", len(chunks)) # Stage 1: Rule-based extraction (only if enabled) rule_results = {} if use_rule_extraction: - print(f" Stage 1: Rule-based entity extraction...") for chunk in chunks: rule_results[chunk.chunk_id] = self.rule_extractor.extract(chunk.content) - # Stage 2: LLM 实体提取 (所有 chunks 并发) - print(f" Stage 2: LLM entity extraction (parallel)...") - llm_results = {} - llm_service = None - + # Stage 2: ONE combined LLM call per chunk (entities + relations) + llm_results: Dict[str, Dict[str, List]] = {} if self.settings.ENABLE_LLM_EXTRACTION: try: - llm_service = await get_llm_service() - llm_results = await self._extract_entities_batch_optimized(chunks) - print(f" LLM entity extraction completed: {len(llm_results)} chunks") + llm_results = await self._extract_entities_and_relations_llm(chunks) + logger.info("LLM combined extraction completed: %d chunks", len(llm_results)) except Exception as e: - print(f" LLM entity extraction failed: {e}") + logger.warning("LLM combined extraction failed: %s", e) - # Stage 3: 合并规则和 LLM 实体 + # Stage 3: 合并规则 + LLM 实体，收集关系 all_entities = [] + all_relations = [] chunk_entities = [] for chunk in chunks: + chunk_result = llm_results.get(chunk.chunk_id, {}) rule_ents = rule_results.get(chunk.chunk_id, []) - llm_ents = llm_results.get(chunk.chunk_id, []) + llm_ents = chunk_result.get("entities", []) merged = self._merge_entity_results(rule_ents, llm_ents) chunk_entities.append({ @@ -453,8 +348,9 @@ async def process_chunks(self, chunks: List[Any], use_rule_extraction: bool = Fa "entities": merged }) all_entities.extend(merged) + all_relations.extend(chunk_result.get("relations", [])) - # Stage 4: 去重 + # Stage 4: 实体去重 entity_dict = {} for entity in all_entities: key = (entity.name.lower(), entity.type) @@ -462,18 +358,10 @@ async def process_chunks(self, chunks: List[Any], use_rule_extraction: bool = Fa entity_dict[key] = entity unique_entities = list(entity_dict.values()) - print(f" Total unique entities: {len(unique_entities)}") - - # Stage 5: LLM 关系提取 (基于合并后的实体) - all_relations = [] - if llm_service and chunk_entities: - print(f" Stage 3: LLM relation extraction (parallel)...") - try: - all_relations = await self._extract_relations_llm_parallel(chunks, chunk_entities, llm_service) - except Exception as e: - print(f" Relation extraction failed: {e}") - - print(f" Total relations extracted: {len(all_relations)}") + logger.info( + "Extraction totals: %d unique entities, %d relations", + len(unique_entities), len(all_relations), + ) return { "entities": unique_entities, diff --git a/backend/app/services/llm.py b/backend/app/services/llm.py index 42a2b67..d4d277f 100644 --- a/backend/app/services/llm.py +++ b/backend/app/services/llm.py @@ -1,12 +1,15 @@ """Bailian (百炼) LLM service.""" import asyncio +import logging import re import httpx -from typing import AsyncGenerator, List, Dict, Any, Optional +from typing import AsyncGenerator, List, Dict, Any, Optional, Tuple import json from app.config import get_settings +logger = logging.getLogger(__name__) + class LLMService: """Service for interacting with Bailian (百炼) API.""" @@ -30,7 +33,7 @@ async def _get_client(self) -> httpx.AsyncClient: async def chat_complete( self, messages: List[Dict[str, str]], - model: str = "qwen-flash", + model: Optional[str] = None, temperature: float = 0.7, max_tokens: int = 8000, stream: bool = False @@ -53,9 +56,10 @@ async def chat_complete( if not self.api_key: raise ValueError("No API key configured") - # Use default model if not specified - if model == "kimi-k2-0905-preview": - model = self.default_model + # Resolve the model from config (settings.BAILIAN_MODEL) when the + # caller doesn't pin one. This is the single source of truth; the old + # literal default in the signature silently overrode the config. + model = model or self.default_model try: response = await client.post( @@ -101,10 +105,11 @@ async def chat_complete( async def chat_complete_stream( self, messages: List[Dict[str, str]], - model: str = "qwen3.5-flash", + model: Optional[str] = None, temperature: float = 0.7, - max_tokens: int = 2000 - ) -> AsyncGenerator[str, None]: + max_tokens: int = 2000, + enable_thinking: Optional[bool] = None, + ) -> AsyncGenerator[Tuple[str, str], None]: """ Stream a chat completion. @@ -113,49 +118,123 @@ async def chat_complete_stream( model: Model to use temperature: Sampling temperature max_tokens: Maximum tokens to generate + enable_thinking: Qwen hybrid-thinking toggle (DashScope + ``enable_thinking``). None = omit the param (legacy + behavior); True = stream the model's reasoning before the + answer; False = answer directly (faster first token). + Non-thinking model tiers reject the param with HTTP 400 — + handled below by dropping it and retrying once, so a stale + toggle degrades to a normal answer instead of an error + bubble. Yields: - Chunks of the generated response + (kind, text) tuples where kind is "content" (answer body) or + "thinking" (reasoning stream — only when thinking is enabled + and the model supports it). """ client = await self._get_client() - # Use default model if not specified - if model == "kimi-k2-0905-preview": - model = self.default_model + # Resolve the model from config (settings.BAILIAN_MODEL) when the + # caller doesn't pin one. This is the single source of truth; the old + # literal default in the signature silently overrode the config. + model = model or self.default_model + + payload: Dict[str, Any] = { + "model": model, + "messages": messages, + "temperature": temperature, + "max_tokens": max_tokens, + "stream": True, + } + if enable_thinking is not None: + payload["enable_thinking"] = enable_thinking + + url = f"{self.base_url}/chat/completions" + headers = { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } try: - async with client.stream( - "POST", - f"{self.base_url}/chat/completions", - headers={ - "Authorization": f"Bearer {self.api_key}", - "Content-Type": "application/json" - }, - json={ - "model": model, - "messages": messages, - "temperature": temperature, - "max_tokens": max_tokens, - "stream": True - } - ) as response: - response.raise_for_status() + async for item in self._stream_completions(client, url, headers, payload): + yield item + except httpx.HTTPStatusError as e: + if ( + enable_thinking is not None + and e.response is not None + and e.response.status_code == 400 + ): + # Non-thinking model tier: the provider rejects the + # enable_thinking param. Drop it and retry ONCE so the user + # gets a normal answer rather than an error bubble. Nothing + # was yielded before raise_for_status fired, so the retry is + # transparent to the caller. + logger.warning( + "enable_thinking=%s rejected by provider (HTTP 400: %.200s); " + "retrying without it", + enable_thinking, e.response.text, + ) + payload.pop("enable_thinking", None) + try: + async for item in self._stream_completions(client, url, headers, payload): + yield item + except Exception as retry_error: + yield ("content", f"\n[Error: {str(retry_error)}]") + else: + yield ("content", f"\n[Error: {str(e)}]") + except Exception as e: + yield ("content", f"\n[Error: {str(e)}]") - async for line in response.aiter_lines(): - if line.startswith("data: "): - data_str = line[6:] - if data_str == "[DONE]": - break - try: - data = json.loads(data_str) - delta = data.get("choices", [{}])[0].get("delta", {}) - if "content" in delta: - yield delta["content"] - except (json.JSONDecodeError, KeyError): - continue + async def _stream_completions( + self, + client: httpx.AsyncClient, + url: str, + headers: Dict[str, str], + payload: Dict[str, Any], + ) -> AsyncGenerator[Tuple[str, str], None]: + """Open one SSE stream and yield (kind, text) delta tuples. + + kind is "thinking" for the ``reasoning_content`` field (Qwen hybrid + thinking) and "content" for the answer body. Raises on HTTP errors + so the caller can decide whether a param-rejection retry applies. + """ + async with client.stream("POST", url, headers=headers, json=payload) as response: + response.raise_for_status() - except Exception as e: - yield f"\n[Error: {str(e)}]" + async for line in response.aiter_lines(): + if line.startswith("data: "): + data_str = line[6:] + if data_str == "[DONE]": + break + try: + data = json.loads(data_str) + # Providers send a trailing usage-only frame with an + # EMPTY choices list; guard before indexing [0] or + # we IndexError (which the outer handler would + # inject into the answer as "[Error: ...]"). + choices = data.get("choices") or [] + if not choices: + continue + delta = choices[0].get("delta", {}) + # Qwen hybrid thinking emits the reasoning stream in + # `reasoning_content` BEFORE the answer body starts. + # Forward it under the "thinking" kind so callers can + # route it to a separate UI block and keep the main + # first-token metric anchored to the answer body. + reasoning = delta.get("reasoning_content") + if reasoning: + yield ("thinking", reasoning) + # OpenAI-compatible providers also emit deltas like + # {"content": null} (role-only / final frames). The + # key is present but the value is None — yielding + # that produced `{"chunk": null}` and later crashed + # the "".join() in the caller. Only forward real, + # non-empty text. + content = delta.get("content") + if content: + yield ("content", content) + except (json.JSONDecodeError, KeyError, IndexError): + continue async def extract_entities_batch( self, @@ -179,8 +258,9 @@ async def extract_entities_batch( Return ONLY a JSON array of objects with format: {{"name": "entity name", "type": "one of {', '.join(entity_types)}", "description": "brief description"}}. If no entities are found, return an empty array.""" - # Limit concurrent requests to avoid 429 - semaphore = asyncio.Semaphore(50) + # Limit concurrent requests to avoid 429 (config-tuned; the old + # hard-coded 50 tripped rate limits whose backoffs slowed the run). + semaphore = asyncio.Semaphore(self.settings.LLM_EXTRACTION_CONCURRENCY) async def _extract_single(text: str) -> List[Dict[str, Any]]: """Extract entities from a single text.""" @@ -191,7 +271,10 @@ async def _extract_single(text: str) -> List[Dict[str, Any]]: ] for attempt in range(3): try: - response = await self.chat_complete(messages, temperature=0.1) + response = await self.chat_complete( + messages, temperature=0.1, + max_tokens=self.settings.LLM_EXTRACT_MAX_TOKENS, + ) json_match = self._extract_json(response) if json_match: entities = json.loads(json_match) @@ -200,13 +283,13 @@ async def _extract_single(text: str) -> List[Dict[str, Any]]: except Exception as e: if "429" in str(e) and attempt < 2: wait_time = (attempt + 1) * 2 - print(f" [LLM Rate Limit] Retrying in {wait_time}s...") + logger.warning("[LLM Rate Limit] Retrying in %ds...", wait_time) await asyncio.sleep(wait_time) continue - print(f" [LLM Entity Extract Error] {e}") + logger.warning("[LLM Entity Extract Error] %s", e) return [] - # Process all texts concurrently (with semaphore limiting to 5) + # Process all texts concurrently (bounded by the semaphore above). tasks = [_extract_single(text) for text in texts] results = await asyncio.gather(*tasks) return list(results) @@ -230,8 +313,8 @@ async def extract_relations_batch( Return ONLY a JSON array of objects with format: {"source": "entity name", "target": "entity name", "relation_type": "relationship type"}. If no relations are found, return an empty array.""" - # Limit concurrent requests - semaphore = asyncio.Semaphore(50) + # Limit concurrent requests (config-tuned, see extract_entities_batch). + semaphore = asyncio.Semaphore(self.settings.LLM_EXTRACTION_CONCURRENCY) async def _extract_single(text: str, entities: List[Dict[str, Any]]) -> List[Dict[str, Any]]: async with semaphore: @@ -253,7 +336,10 @@ async def _extract_single(text: str, entities: List[Dict[str, Any]]) -> List[Dic for attempt in range(3): try: - response = await self.chat_complete(messages, temperature=0.1) + response = await self.chat_complete( + messages, temperature=0.1, + max_tokens=self.settings.LLM_EXTRACT_MAX_TOKENS, + ) json_match = self._extract_json(response) if json_match: relations = json.loads(json_match) @@ -262,10 +348,10 @@ async def _extract_single(text: str, entities: List[Dict[str, Any]]) -> List[Dic except Exception as e: if "429" in str(e) and attempt < 2: wait_time = (attempt + 1) * 2 - print(f" [LLM Rate Limit] Retrying in {wait_time}s...") + logger.warning("[LLM Rate Limit] Retrying in %ds...", wait_time) await asyncio.sleep(wait_time) continue - print(f" [LLM Relation Extract Error] {e}") + logger.warning("[LLM Relation Extract Error] %s", e) return [] # Process concurrently @@ -273,6 +359,70 @@ async def _extract_single(text: str, entities: List[Dict[str, Any]]) -> List[Dic results = await asyncio.gather(*tasks) return list(results) + async def extract_entities_and_relations_batch( + self, + texts: List[str], + entity_types: List[str] = None, + ) -> List[Dict[str, List[Dict[str, Any]]]]: + """Extract entities AND relations from multiple texts, ONE LLM call per text. + + Replaces the old two-stage pipeline (extract_entities_batch, then a + full stage barrier, then extract_relations_batch): that design paid + 2N round-trips for N texts, and no relation call could start until + EVERY chunk's entities had come back. Here each text makes a single + call returning both, so LLM work is halved and fully overlapped. + + Args: + texts: List of texts to process. + entity_types: Allowed entity types named in the prompt. + + Returns: + One dict per input text (same order, same length): + {"entities": [{"name","type","description"}, ...], + "relations": [{"source","target","relation_type"}, ...]} + A failed or unparseable text yields empty arrays — extraction is + best-effort and must never take down the whole document. + """ + if entity_types is None: + entity_types = ["PERSON", "ORGANIZATION", "LOCATION", "CONCEPT", "EVENT"] + + system_prompt = f"""You are a knowledge extraction assistant. Extract entities AND the relations between them from the given text. +Return ONLY a JSON object with exactly this shape: +{{"entities": [{{"name": "entity name", "type": "one of {', '.join(entity_types)}", "description": "brief description"}}], + "relations": [{{"source": "entity name", "target": "entity name", "relation_type": "short relation label"}}]}} +Rules: +- Every relation's "source" and "target" MUST be names that appear in the "entities" array. +- If nothing is found, return {{"entities": [], "relations": []}}.""" + + semaphore = asyncio.Semaphore(self.settings.LLM_EXTRACTION_CONCURRENCY) + + async def _extract_single(text: str) -> Dict[str, List[Dict[str, Any]]]: + async with semaphore: + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": f"Extract entities and relations from:\n\n{text[:2000]}"}, + ] + for attempt in range(3): + try: + response = await self.chat_complete( + messages, temperature=0.1, + max_tokens=self.settings.LLM_EXTRACT_MAX_TOKENS, + ) + return self._parse_extraction_response(response) + except Exception as e: + if "429" in str(e) and attempt < 2: + wait_time = (attempt + 1) * 2 + logger.warning("[LLM Rate Limit] Retrying in %ds...", wait_time) + await asyncio.sleep(wait_time) + continue + logger.warning("[LLM Combined Extract Error] %s", e) + return {"entities": [], "relations": []} + return {"entities": [], "relations": []} + + tasks = [_extract_single(text) for text in texts] + results = await asyncio.gather(*tasks) + return list(results) + async def generate_rag_response( self, query: str, @@ -365,71 +515,79 @@ def _extract_json(self, text: str) -> Optional[str]: return match.group(0) return None - async def generate_followups( - self, - query: str, - answer: str, - n: int = 3, - ) -> List[str]: - """Generate up to `n` follow-up question chips based on the just- - given answer. The returned list is capped at n, blank entries - are dropped, and the call NEVER raises — a failure here is a - UX bonus, not a blocking dependency. - - Returns an empty list if the LLM call fails, returns non-JSON, - or returns a list that has no usable strings after cleaning. + def _extract_json_object(self, text: str) -> Optional[str]: + """Extract the first JSON object ({...}) from text. + + Greedy match to the LAST '}' — correct because the extraction + prompts demand "Return ONLY a JSON object", so the object is the + whole payload and there is no trailing prose whose braces could + over-extend the match. """ - system_prompt = ( - "You are a follow-up question generator. Given a user's " - "question and the assistant's answer, suggest the next " - f"{n} questions the user is most likely to ask to go " - "deeper. Output ONLY a JSON array of strings — no prose, " - "no markdown fences, no numbering. Each entry should be a " - "complete, natural-language question. Example shape: " - '["How does X relate to Y?", "What about Z?", "Why does W?"]' - ) - user_prompt = ( - f"User question:\n{query}\n\n" - f"Assistant answer:\n{(answer or '')[:2000]}\n\n" - f"Return the {n} follow-up questions." - ) - try: - raw = await self.chat_complete( - [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt}, - ], - temperature=0.7, - max_tokens=400, - ) - except Exception as e: - print(f" [LLM Followups Error] {e}") - return [] - - # The LLM may have wrapped the array in ```json … ``` fences or - # a prose preamble — `_extract_json` finds the first JSON array - # in the text. If there's none at all we bail out. - json_str = self._extract_json(raw or "") - if not json_str: - return [] - try: - parsed = json.loads(json_str) - except (json.JSONDecodeError, ValueError): - return [] - if not isinstance(parsed, list): - return [] - - # Cap at n, drop non-strings, strip whitespace, drop empties. - cleaned: List[str] = [] - for item in parsed: - if not isinstance(item, str): - continue - s = item.strip() - if s: - cleaned.append(s) - if len(cleaned) >= n: - break - return cleaned + match = re.search(r'\{.*\}', text, re.DOTALL) + if match: + return match.group(0) + return None + + def _parse_extraction_response( + self, text: str + ) -> Dict[str, List[Dict[str, Any]]]: + """Normalise a combined entity+relation LLM response. + + Accepts the expected object shape {"entities": [...], "relations": + [...]} and, as a fallback, a bare entity array (older prompt shape + some models revert to). Always returns both keys with list values; + entries that are not dicts or lack the required name / source+target + fields are dropped, so downstream code never has to re-validate. + """ + empty: Dict[str, List[Dict[str, Any]]] = {"entities": [], "relations": []} + + # Pick the parse path by whichever structural token appears first: + # '{' → expected object shape; '[' → legacy bare-entity-array shape. + # Position matters — blindly trying the object regex first would + # hijack the INNER {...} of a bare array and discard real entities. + first_obj = text.find("{") + first_arr = text.find("[") + + if first_obj != -1 and (first_arr == -1 or first_obj < first_arr): + obj_str = self._extract_json_object(text) + if obj_str: + try: + parsed = json.loads(obj_str) + except (json.JSONDecodeError, ValueError): + parsed = None + if isinstance(parsed, dict): + entities = parsed.get("entities") + relations = parsed.get("relations") + return { + "entities": [ + e for e in entities + if isinstance(e, dict) and str(e.get("name") or "").strip() + ] if isinstance(entities, list) else [], + "relations": [ + r for r in relations + if isinstance(r, dict) + and str(r.get("source") or "").strip() + and str(r.get("target") or "").strip() + ] if isinstance(relations, list) else [], + } + + # Fallback: bare array → treat as entities-only rather than + # discarding the work. + arr_str = self._extract_json(text) + if arr_str: + try: + parsed = json.loads(arr_str) + except (json.JSONDecodeError, ValueError): + return empty + if isinstance(parsed, list): + return { + "entities": [ + e for e in parsed + if isinstance(e, dict) and str(e.get("name") or "").strip() + ], + "relations": [], + } + return empty async def close(self): """Close HTTP client.""" @@ -448,3 +606,11 @@ async def get_llm_service() -> LLMService: if _llm_service is None: _llm_service = LLMService() return _llm_service + + +async def close_llm_service() -> None: + """Close the shared LLM HTTP client at shutdown (no-op if never created).""" + global _llm_service + if _llm_service is not None: + await _llm_service.close() + _llm_service = None diff --git a/backend/app/services/neo4j_client.py b/backend/app/services/neo4j_client.py index 6c83235..01a817b 100644 --- a/backend/app/services/neo4j_client.py +++ b/backend/app/services/neo4j_client.py @@ -113,6 +113,66 @@ async def create_chunk_links(self, chunk_id: str, prev_chunk_id: Optional[str], MERGE (curr)-[:NEXT]->(next) """, chunk_id=chunk_id, next_id=next_chunk_id) + async def create_chunk_nodes_batch( + self, doc_id: str, user_id: int, chunks: List[Dict[str, Any]] + ) -> int: + """Bulk-create Chunk nodes + CONTAINS links in ONE round-trip (UNWIND). + + Replaces the per-chunk create_chunk_node loop in the ingestion + pipeline: an N-chunk document used to cost 2N+ serial Neo4j + round-trips (~5–15ms each); this is a single transaction. + + Each chunk dict must contain: chunk_id, content, hierarchy_path + (list[str]), position (int). Node properties match create_chunk_node + exactly, so downstream queries (CONTAINS-edge traversals, graph + visualisation) see no difference. + + The Document node is MERGEd once before the UNWIND (the pipeline + always creates it earlier via create_document_node; the MERGE keeps + this helper safe to call standalone, mirroring create_chunk_node's + old per-call MERGE). + """ + if not chunks: + return 0 + async with self.session() as session: + result = await session.run(""" + MERGE (d:Document {doc_id: $doc_id}) + SET d.user_id = $user_id + WITH d + UNWIND $chunks AS ch + MERGE (c:Chunk {chunk_id: ch.chunk_id}) + SET c.content = ch.content, + c.hierarchy_path = ch.hierarchy_path, + c.position = ch.position, + c.user_id = $user_id, + c.created_at = datetime() + MERGE (d)-[:CONTAINS]->(c) + RETURN count(c) AS created + """, doc_id=doc_id, user_id=user_id, chunks=chunks) + record = await result.single() + return record["created"] if record else 0 + + async def create_chunk_links_batch(self, links: List[Dict[str, str]]) -> int: + """Bulk-create NEXT links between chunks in ONE round-trip (UNWIND). + + Each link dict must contain: from_id, to_id. The caller is expected + to have deduplicated pairs already (prev/next pointers are + symmetric, so naive expansion double-counts every edge — MERGE + would dedupe anyway, but sending unique pairs halves the work). + """ + if not links: + return 0 + async with self.session() as session: + result = await session.run(""" + UNWIND $links AS link + MATCH (a:Chunk {chunk_id: link.from_id}) + MATCH (b:Chunk {chunk_id: link.to_id}) + MERGE (a)-[:NEXT]->(b) + RETURN count(*) AS linked + """, links=links) + record = await result.single() + return record["linked"] if record else 0 + async def create_entity(self, name: str, entity_type: str, description: Optional[str], user_id: int) -> str: """Create or update an entity node.""" @@ -419,7 +479,7 @@ async def create_relations_batch( async def search_entities(self, query: str, user_id: int, limit: int = 10) -> List[Dict[str, Any]]: """Search entities by name (case-insensitive).""" - print(f" Neo4j search: query='{query}', user_id={user_id}") + logger.debug(f" Neo4j search: query='{query}', user_id={user_id}") async with self.session() as session: # First check all entities for this user check_result = await session.run(""" @@ -429,7 +489,7 @@ async def search_entities(self, query: str, user_id: int, limit: int = 10) -> Li """, user_id=user_id) check_record = await check_result.single() total_entities = check_record["total"] if check_record else 0 - print(f" Total entities for user {user_id}: {total_entities}") + logger.debug(f" Total entities for user {user_id}: {total_entities}") # Now search result = await session.run(""" @@ -439,7 +499,7 @@ async def search_entities(self, query: str, user_id: int, limit: int = 10) -> Li LIMIT $limit """, search_term=query, user_id=user_id, limit=limit) entities = [record.data() async for record in result] - print(f" Found {len(entities)} matching entities") + logger.debug(f" Found {len(entities)} matching entities") return entities async def get_related_entities(self, entity_names: List[str], user_id: int, @@ -930,7 +990,7 @@ async def get_entity_graph_for_visualization(self, query: str, user_id: int, async def get_full_graph_for_visualization(self, user_id: int) -> Dict[str, Any]: """Get complete graph with ALL nodes and relationships for visualization.""" - print(f" Getting full graph for user {user_id}") + logger.debug(f" Getting full graph for user {user_id}") async with self.session() as session: nodes = {} edges = [] @@ -1010,7 +1070,7 @@ async def get_full_graph_for_visualization(self, user_id: int) -> Dict[str, Any] "properties": {"user_id": record['user_id']} } - print(f" Found {len(nodes)} total nodes") + logger.debug(f" Found {len(nodes)} total nodes") # Get all relationships between these nodes # OWNS: (User)-[:OWNS]->(Document) @@ -1098,7 +1158,7 @@ async def get_full_graph_for_visualization(self, user_id: int) -> Dict[str, Any] "type": "RELATES_TO" }) - print(f" Found {len(edges)} total relationships") + logger.debug(f" Found {len(edges)} total relationships") return { "nodes": list(nodes.values()), @@ -1107,7 +1167,7 @@ async def get_full_graph_for_visualization(self, user_id: int) -> Dict[str, Any] async def delete_user_data(self, user_id: int): """Delete all data for a user.""" - print(f"[neo4j] Deleting all data for user {user_id}") + logger.debug(f"[neo4j] Deleting all data for user {user_id}") async with self.session() as session: # Delete all relations first await session.run(""" @@ -1137,11 +1197,11 @@ async def delete_user_data(self, user_id: int): DETACH DELETE d """, user_id=user_id) - print(f"[neo4j] All data deleted for user {user_id}") + logger.debug(f"[neo4j] All data deleted for user {user_id}") async def delete_document(self, doc_id: str, user_id: int): """Delete a document and its chunks, entities, and relations.""" - print(f"[neo4j] ====== DELETE START: doc_id={doc_id}, user_id={user_id} ======") + logger.debug(f"[neo4j] ====== DELETE START: doc_id={doc_id}, user_id={user_id} ======") async with self.session() as session: # First, let's see what's currently in the database @@ -1151,7 +1211,7 @@ async def delete_document(self, doc_id: str, user_id: int): RETURN count(e) as count """, user_id=user_id) record = await result.single() - print(f"[neo4j] BEFORE DELETE: Total entities in DB: {record['count'] if record else 0}") + logger.debug(f"[neo4j] BEFORE DELETE: Total entities in DB: {record['count'] if record else 0}") # Check how many chunks exist for this user result = await session.run(""" @@ -1160,7 +1220,7 @@ async def delete_document(self, doc_id: str, user_id: int): RETURN count(c) as count """, user_id=user_id) record = await result.single() - print(f"[neo4j] BEFORE DELETE: Total chunks in DB: {record['count'] if record else 0}") + logger.debug(f"[neo4j] BEFORE DELETE: Total chunks in DB: {record['count'] if record else 0}") # Step 1: Check if Document exists, collect chunk IDs result = await session.run(""" @@ -1173,7 +1233,7 @@ async def delete_document(self, doc_id: str, user_id: int): chunk_ids = record["chunk_ids"] if record else [] if not doc_exists: - print(f"[neo4j] Document {doc_id} not found in Neo4j!") + logger.debug(f"[neo4j] Document {doc_id} not found in Neo4j!") # Still print stats result = await session.run(""" MATCH (e:Entity) @@ -1181,10 +1241,10 @@ async def delete_document(self, doc_id: str, user_id: int): RETURN count(e) as count """, user_id=user_id) record = await result.single() - print(f"[neo4j] AFTER DELETE: Total entities: {record['count'] if record else 0}") + logger.debug(f"[neo4j] AFTER DELETE: Total entities: {record['count'] if record else 0}") return else: - print(f"[neo4j] Found document with {len(chunk_ids)} chunks") + logger.debug(f"[neo4j] Found document with {len(chunk_ids)} chunks") # Step 2: Collect entity names that are mentioned in this document's chunks result = await session.run(""" @@ -1194,7 +1254,7 @@ async def delete_document(self, doc_id: str, user_id: int): """, chunk_ids=chunk_ids, user_id=user_id) record = await result.single() entity_names = record["entity_names"] if record else [] - print(f"[neo4j] Entities in THIS document: {len(entity_names)}") + logger.debug(f"[neo4j] Entities in THIS document: {len(entity_names)}") # Step 3: Delete MENTIONS relations from chunks if chunk_ids: @@ -1205,7 +1265,7 @@ async def delete_document(self, doc_id: str, user_id: int): RETURN count(r) as deleted """, chunk_ids=chunk_ids) record = await result.single() - print(f"[neo4j] Step 3: Deleted {record['deleted'] if record else 0} MENTIONS relations") + logger.debug(f"[neo4j] Step 3: Deleted {record['deleted'] if record else 0} MENTIONS relations") # Step 4: Delete entities that were IN THIS DOCUMENT only if no # remaining chunk (of the SAME user) still mentions them. The @@ -1236,7 +1296,7 @@ async def delete_document(self, doc_id: str, user_id: int): RETURN count(r) as deleted """, user_id=user_id, entity_names=entity_names) record = await result.single() - print(f"[neo4j] Step 5: Deleted {record['deleted'] if record else 0} RELATES_TO relations") + logger.debug(f"[neo4j] Step 5: Deleted {record['deleted'] if record else 0} RELATES_TO relations") # Step 6: Delete Document and Chunk nodes result = await session.run(""" @@ -1246,7 +1306,7 @@ async def delete_document(self, doc_id: str, user_id: int): RETURN count(d) as deleted """, doc_id=doc_id) record = await result.single() - print(f"[neo4j] Step 6: Deleted {record['deleted'] if record else 0} documents") + logger.debug(f"[neo4j] Step 6: Deleted {record['deleted'] if record else 0} documents") # Final stats result = await session.run(""" @@ -1255,7 +1315,7 @@ async def delete_document(self, doc_id: str, user_id: int): RETURN count(e) as count """, user_id=user_id) record = await result.single() - print(f"[neo4j] ====== DELETE COMPLETE: Remaining entities: {record['count'] if record else 0} ======") + logger.debug(f"[neo4j] ====== DELETE COMPLETE: Remaining entities: {record['count'] if record else 0} ======") # Singleton instance diff --git a/backend/app/services/progress_tracker.py b/backend/app/services/progress_tracker.py index d73a290..41fe4af 100644 --- a/backend/app/services/progress_tracker.py +++ b/backend/app/services/progress_tracker.py @@ -1,43 +1,59 @@ """Progress tracking using Server-Sent Events (SSE).""" import asyncio import json -from typing import Dict, Callable, Optional, List +import logging +from typing import Dict, Callable, Optional, List, Set from collections import defaultdict import aiosqlite from app.config import get_settings +logger = logging.getLogger(__name__) + class ProgressEmitter: """Centralized progress event emitter using asyncio.""" def __init__(self): - self._subscribers: Dict[str, asyncio.Queue] = defaultdict(asyncio.Queue) - self._locks: Dict[str, asyncio.Lock] = {} + # Each subscriber gets its OWN queue. A single shared queue per doc_id + # made multiple watchers (e.g. two browser tabs) steal events + # round-robin, and any one client disconnecting deleted the queue out + # from under the others. A set of per-subscriber queues fixes both. + self._subscribers: Dict[str, Set[asyncio.Queue]] = defaultdict(set) def subscribe(self, doc_id: str) -> asyncio.Queue: - """Subscribe to progress updates for a document.""" - if doc_id not in self._locks: - self._locks[doc_id] = asyncio.Lock() - return self._subscribers[doc_id] - - def unsubscribe(self, doc_id: str): - """Unsubscribe from progress updates.""" - if doc_id in self._subscribers: - del self._subscribers[doc_id] - if doc_id in self._locks: - del self._locks[doc_id] + """Create and register a fresh per-subscriber queue for a document.""" + queue: asyncio.Queue = asyncio.Queue() + self._subscribers[doc_id].add(queue) + return queue + + def unsubscribe(self, doc_id: str, queue: Optional[asyncio.Queue] = None): + """Remove a subscriber's queue. + + When ``queue`` is given, only that subscriber is removed (preferred). + Omitting it drops every subscriber for the doc — kept for backwards + compatibility, but avoid it when multiple clients may be watching. + """ + if queue is None: + self._subscribers.pop(doc_id, None) + return + queues = self._subscribers.get(doc_id) + if queues is not None: + queues.discard(queue) + if not queues: + self._subscribers.pop(doc_id, None) async def emit(self, doc_id: str, progress_type: str, message: str, data: dict = None): - """Emit a progress event.""" + """Emit a progress event to every subscriber of a document.""" event = { "type": progress_type, "message": message, "data": data or {} } - if doc_id in self._subscribers: - await self._subscribers[doc_id].put(event) + # Iterate a snapshot so a concurrent unsubscribe can't mutate the set. + for queue in list(self._subscribers.get(doc_id, ())): + await queue.put(event) async def emit_and_save(self, doc_id: str, user_id: int, progress_type: str, message: str, data: dict = None, entity_count: int = 0, relation_count: int = 0): diff --git a/backend/app/services/reranker.py b/backend/app/services/reranker.py index 4dfee7f..f7bf4cd 100644 --- a/backend/app/services/reranker.py +++ b/backend/app/services/reranker.py @@ -89,3 +89,11 @@ async def get_rerank_service() -> RerankService: if _rerank_service is None: _rerank_service = RerankService() return _rerank_service + + +async def close_rerank_service() -> None: + """Close the shared rerank HTTP client at shutdown (no-op if never created).""" + global _rerank_service + if _rerank_service is not None: + await _rerank_service.close() + _rerank_service = None diff --git a/backend/eval/gold/06_followup.json b/backend/eval/gold/06_followup.json deleted file mode 100644 index 032bbbb..0000000 --- a/backend/eval/gold/06_followup.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_comment": [ - "Follow-up suggestions question. The user asks anything; we expect", - "the backend to ALSO return up to 3 follow-up question chips after", - "the answer. The LLM service is unit-tested separately for the", - "parsing/cleaning logic (see tests/test_followups.py); this gold case", - "verifies the end-to-end envelope: the response (or SSE event) must", - "carry a 'followups' field that's a list of 1-3 natural-language", - "questions related to the original query and answer.", - "", - "Use with with_followups=True (the default).", - "expected_keywords are NOT for the answer — they're for the", - "followups themselves: a couple of phrases that prove the LLM", - "understood the topic and is suggesting next-step questions rather", - "than rephrasing the original." - ], - "id": "followup-01", - "query": "Explain how QLoRA works.", - "with_followups": true, - "expected_keywords": [ - "LoRA", - "memory" - ], - "min_followups": 1, - "difficulty": "easy", - "tags": ["ux", "followup"] -} diff --git a/backend/scripts/rebuild_chroma.py b/backend/scripts/rebuild_chroma.py new file mode 100644 index 0000000..61c1a82 --- /dev/null +++ b/backend/scripts/rebuild_chroma.py @@ -0,0 +1,139 @@ +"""One-off: rebuild Chroma collection from SQLite + embedding cache. + +Vectors were lost when the external Chroma container (pure in-memory) was +stopped. The chunk metadata survives in the SQLite ``chunks`` table, and +every embedding is cached in ``embedding_cache`` keyed by md5(content) + +model — so we can repopulate Chroma with ZERO LLM / embedding API calls. + +Reuses ``get_chroma_client()`` so the collection name +(``knowledge_graph_chunks``), cosine distance, and upsert semantics match +the ingestion path exactly. Metadata fields mirror ``documents.py`` lines +287-296 verbatim, so downstream ``where`` filters (user_id, document_id) +and ``get_chunk_context`` (prev/next_chunk_id) keep working. + +Run from backend/ (so .env is picked up by pydantic-settings): + + .venv/Scripts/python scripts/rebuild_chroma.py +""" +from __future__ import annotations + +import hashlib +import json +import sqlite3 +import sys +from pathlib import Path +from typing import Any, List, Tuple + +# Allow `from app...` when run as a plain script. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from app.config import get_settings # noqa: E402 +from app.services.chroma_client import get_chroma_client # noqa: E402 + +BATCH = 100 + + +def _text_hash(text: str) -> str: + """md5 hex of UTF-8 text — must match EmbeddingService._get_text_hash.""" + return hashlib.md5(text.encode("utf-8")).hexdigest() + + +def _load_chunks(db_path: Path, model: str) -> Tuple[List[dict], List[str]]: + """Read every chunk row and attach its cached embedding. + + Returns (rows_with_embedding, chunk_ids_missing_cache). A chunk is + skipped (not silently dropped) when its content is blank or no cached + embedding exists — those need an API re-embed and are reported. + """ + conn = sqlite3.connect(str(db_path)) + conn.row_factory = sqlite3.Row + cur = conn.execute( + "SELECT chunk_id, document_id, user_id, content, hierarchy_path, " + "level, prev_chunk_id, next_chunk_id FROM chunks" + ) + rows = cur.fetchall() + + ready: List[dict] = [] + missing: List[str] = [] + for r in rows: + content = r["content"] or "" + if not content.strip(): + print(f"[warn] blank content, skipping chunk_id={r['chunk_id']}") + continue + cached = conn.execute( + "SELECT embedding FROM embedding_cache " + "WHERE text_hash = ? AND model = ?", + (_text_hash(content), model), + ).fetchone() + if not cached: + missing.append(r["chunk_id"]) + continue + embedding: List[float] = json.loads(cached["embedding"].decode("utf-8")) + + # hierarchy_path is stored in SQLite as comma-no-space; the ingestion + # path writes ", " (comma-space) into Chroma. Reconstruct so a future + # re-ingest is byte-identical and idempotent. + hp_raw = r["hierarchy_path"] or "" + hierarchy_path = ", ".join(p for p in hp_raw.split(",") if p) + metadata: dict[str, Any] = { + "user_id": str(r["user_id"]), + "document_id": r["document_id"], + "hierarchy_level": str(r["level"]) if r["level"] is not None else "", + "hierarchy_path": hierarchy_path, + "prev_chunk_id": r["prev_chunk_id"] or "", + "next_chunk_id": r["next_chunk_id"] or "", + } + ready.append({ + "chunk_id": r["chunk_id"], + "content": content, + "embedding": embedding, + "metadata": metadata, + }) + conn.close() + return ready, missing + + +def main() -> int: + settings = get_settings() + + db_path = Path(settings.SQLITE_PATH) + if not db_path.is_absolute(): + db_path = Path(__file__).resolve().parent.parent / db_path + if not db_path.exists(): + print(f"[error] SQLite DB not found at {db_path}") + return 1 + print(f"[info] SQLite: {db_path}") + + ready, missing = _load_chunks(db_path, settings.EMBEDDING_MODEL) + print(f"[info] chunks with cached embedding: {len(ready)}") + if missing: + print(f"[warn] cache miss for {len(missing)} chunk(s) — API re-embed needed:") + for cid in missing[:10]: + print(f" {cid}") + if len(missing) > 10: + print(f" ...and {len(missing) - 10} more") + + if not ready: + print("[error] nothing to upsert; aborting.") + return 1 + + chroma = get_chroma_client() + total = 0 + for i in range(0, len(ready), BATCH): + batch = ready[i:i + BATCH] + chroma.add_chunks( + chunk_ids=[b["chunk_id"] for b in batch], + documents=[b["content"] for b in batch], + embeddings=[b["embedding"] for b in batch], + metadatas=[b["metadata"] for b in batch], + ) + total += len(batch) + print(f"[info] upserted batch {i // BATCH + 1}: {total}/{len(ready)}") + + count = chroma._collection.count() + print(f"[done] upserted {total} chunks; collection now holds {count} vectors") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py new file mode 100644 index 0000000..d336dd9 --- /dev/null +++ b/backend/tests/conftest.py @@ -0,0 +1,15 @@ +"""Pytest configuration. + +``get_settings()`` refuses to start with a known placeholder JWT_SECRET (a +security fix). The test suite therefore needs a strong, throwaway value set in +the real environment BEFORE any app module is imported or the settings cache is +populated. pytest imports conftest.py first, so setting it here is safe. + +Environment variables take precedence over the ``.env`` file in +pydantic-settings, so this also overrides any placeholder that may sit in +``backend/.env``. +""" +import os + +os.environ.setdefault("JWT_SECRET", "test-only-" + "k" * 48) +os.environ.setdefault("APP_ENV", "test") diff --git a/backend/tests/test_doc_status.py b/backend/tests/test_doc_status.py new file mode 100644 index 0000000..262db68 --- /dev/null +++ b/backend/tests/test_doc_status.py @@ -0,0 +1,219 @@ +"""Tests for the document processing state machine (services/doc_status.py). + +Covers transition rules, idempotency, failure recording, terminal guards, +retry reset, the legacy-schema migration, and the fresh-schema default. +""" +import asyncio +import os + +import aiosqlite +import pytest + +from app.config import get_settings +from app.database import init_db, get_db +from app.services.doc_status import ( + DocStatus, + DocumentNotFound, + InvalidStatusTransition, + get_document_status, + reset_for_retry, + set_document_status, + validate_transition, +) + + +@pytest.fixture(autouse=True) +def tmp_sqlite(monkeypatch, tmp_path): + """Give every test a throwaway SQLite file.""" + monkeypatch.setenv("SQLITE_PATH", str(tmp_path / "doc_status_test.db")) + get_settings.cache_clear() + yield + get_settings.cache_clear() + + +async def _bootstrap(user_id: int = 1) -> None: + await init_db() + async with get_db() as db: + await db.execute( + "INSERT OR IGNORE INTO users (id, username, password_hash) " + "VALUES (?, ?, ?)", + (user_id, f"user{user_id}", "x"), + ) + await db.commit() + + +async def _insert_doc(doc_id: str, status: str = "pending", user_id: int = 1) -> None: + async with get_db() as db: + await db.execute( + "INSERT INTO documents (id, user_id, title, status) " + "VALUES (?, ?, ?, ?)", + (doc_id, user_id, "t", status), + ) + await db.commit() + + +# --------------------------------------------------------------------------- +# Pure transition rules (no DB) +# --------------------------------------------------------------------------- +def test_validate_transition_rules(): + assert validate_transition("pending", "document_created") is True + assert validate_transition("pending", "indexed") is True # skip-ahead ok + assert validate_transition("indexed", "indexed") is False # idempotent + assert validate_transition("graphed", "document_created") is False # backward no-op + assert validate_transition("indexed", "failed") is True # may fail + with pytest.raises(InvalidStatusTransition): + validate_transition("ready", "failed") # terminal + with pytest.raises(InvalidStatusTransition): + validate_transition("failed", "indexed") # terminal + + +# --------------------------------------------------------------------------- +# DB-backed behaviour +# --------------------------------------------------------------------------- +def test_forward_transitions_apply(): + async def main(): + await _bootstrap() + await _insert_doc("d1") + assert await set_document_status("d1", DocStatus.DOCUMENT_CREATED) is True + assert await get_document_status("d1") == "document_created" + assert await set_document_status("d1", DocStatus.INDEXED) is True + assert await set_document_status("d1", DocStatus.GRAPHED) is True + assert await set_document_status("d1", DocStatus.READY) is True + assert await get_document_status("d1") == "ready" + + asyncio.run(main()) + + +def test_same_and_backward_are_noops(): + async def main(): + await _bootstrap() + await _insert_doc("d2", status="indexed") + assert await set_document_status("d2", DocStatus.INDEXED) is False + assert await set_document_status("d2", DocStatus.DOCUMENT_CREATED) is False + assert await get_document_status("d2") == "indexed" + + asyncio.run(main()) + + +def test_failed_records_error_and_blocks_forward(): + async def main(): + await _bootstrap() + await _insert_doc("d3", status="graphed") + assert await set_document_status( + "d3", DocStatus.FAILED, error_message="boom" + ) is True + async with get_db() as db: + cur = await db.execute( + "SELECT status, error_message FROM documents WHERE id = 'd3'" + ) + row = await cur.fetchone() + assert row["status"] == "failed" + assert row["error_message"] == "boom" + # Cannot leave 'failed' without an explicit reset. + with pytest.raises(InvalidStatusTransition): + await set_document_status("d3", DocStatus.READY) + # Re-failing is an idempotent no-op. + assert await set_document_status("d3", DocStatus.FAILED) is False + + asyncio.run(main()) + + +def test_successful_transition_clears_error_message(): + async def main(): + await _bootstrap() + await _insert_doc("dc", status="pending") + # fail it, then reset + advance; error_message must be cleared. + await set_document_status("dc", DocStatus.FAILED, error_message="x") + await reset_for_retry("dc") + await set_document_status("dc", DocStatus.INDEXED) + async with get_db() as db: + cur = await db.execute( + "SELECT status, error_message FROM documents WHERE id = 'dc'" + ) + row = await cur.fetchone() + assert row["status"] == "indexed" + assert row["error_message"] is None + + asyncio.run(main()) + + +def test_ready_is_terminal(): + async def main(): + await _bootstrap() + await _insert_doc("d4", status="ready") + with pytest.raises(InvalidStatusTransition): + await set_document_status("d4", DocStatus.FAILED) + + asyncio.run(main()) + + +def test_reset_for_retry(): + async def main(): + await _bootstrap() + await _insert_doc("d5", status="failed") + assert await reset_for_retry("d5") is True + assert await get_document_status("d5") == "pending" + # Reset from a non-failed state is rejected. + await _insert_doc("d6", status="indexed") + with pytest.raises(InvalidStatusTransition): + await reset_for_retry("d6") + + asyncio.run(main()) + + +def test_missing_document_raises(): + async def main(): + await _bootstrap() + with pytest.raises(DocumentNotFound): + await set_document_status("nope", DocStatus.READY) + assert await get_document_status("nope") is None + + asyncio.run(main()) + + +def test_new_document_defaults_to_pending(): + async def main(): + await _bootstrap() + async with get_db() as db: + await db.execute( + "INSERT INTO documents (id, user_id, title) VALUES ('dflt', 1, 't')" + ) + await db.commit() + assert await get_document_status("dflt") == "pending" + + asyncio.run(main()) + + +def test_migration_backfills_legacy_rows(): + """A pre-state-machine documents table gains the columns; old rows -> ready.""" + + async def main(): + settings = get_settings() + os.makedirs(os.path.dirname(settings.SQLITE_PATH), exist_ok=True) + # Legacy schema: no status / error_message / updated_at columns. + async with aiosqlite.connect(settings.SQLITE_PATH) as db: + await db.execute( + "CREATE TABLE documents (" + "id TEXT PRIMARY KEY, user_id INTEGER, title TEXT, " + "file_path TEXT, original_filename TEXT, file_type TEXT, " + "created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)" + ) + await db.execute( + "INSERT INTO documents (id, user_id, title) " + "VALUES ('legacy', 1, 'old')" + ) + await db.commit() + + await init_db() # should ALTER + backfill + + async with get_db() as db: + cur = await db.execute( + "SELECT status, error_message, updated_at " + "FROM documents WHERE id = 'legacy'" + ) + row = await cur.fetchone() + assert row["status"] == "ready" + assert row["error_message"] is None + assert row["updated_at"] is not None + + asyncio.run(main()) diff --git a/backend/tests/test_document_detail.py b/backend/tests/test_document_detail.py index f1a58e6..bac9a07 100644 --- a/backend/tests/test_document_detail.py +++ b/backend/tests/test_document_detail.py @@ -334,6 +334,7 @@ def test_detail_endpoint_returns_combined_payload(): doc_row = _FakeRow({ "id": "d-1", "title": "Hello", "original_filename": "hello.pdf", "file_type": "pdf", "file_size": 1024, "created_at": "2026-01-01", + "status": "ready", "error_message": None, }) tags_rows = [ _FakeRow({"tag": "research"}), @@ -424,6 +425,7 @@ def test_detail_endpoint_user_scoped(): doc_row = _FakeRow({ "id": "d-1", "title": "Private", "original_filename": "p.pdf", "file_type": "pdf", "file_size": 100, "created_at": "2026-01-01", + "status": "ready", "error_message": None, }) # Doc is owned by user 1, but we'll authenticate as user 999. # The fake mimics the WHERE user_id=? filter: anyone else's id diff --git a/backend/tests/test_embedding.py b/backend/tests/test_embedding.py index 82cd05a..23ace35 100644 --- a/backend/tests/test_embedding.py +++ b/backend/tests/test_embedding.py @@ -95,6 +95,7 @@ def make_service(): svc.api_key = svc.settings.SILICON_FLOW_API_KEY svc.model = svc.settings.EMBEDDING_MODEL svc._semaphore = asyncio.Semaphore(5) + svc._client = None # lazy shared client, created on first _get_client() return svc diff --git a/backend/tests/test_followups.py b/backend/tests/test_followups.py deleted file mode 100644 index 5cc679f..0000000 --- a/backend/tests/test_followups.py +++ /dev/null @@ -1,358 +0,0 @@ -"""Tests for #0 follow-up suggestions. - -Standalone runner: - cd backend - ../.venv/Scripts/python.exe tests/test_followups.py - -Coverage: - 1. ChatRequest.with_followups defaults to True - 2. ChatRequest.with_followups=False is honored - 3. ChatResponse.followups field is present and defaults to [] - 4. LLMService.generate_followups method exists with correct signature - 5. generate_followups parses a clean JSON array of 3 questions - 6. generate_followups parses JSON with surrounding prose (extract_json) - 7. generate_followups caps at 3 results and drops blanks - 8. generate_followups returns [] on LLM error (never raises to caller) - 9. generate_followups returns [] when the JSON has fewer than 3 entries - 10. Stream generator emits a 'followups' SSE event when with_followups=True - 11. Stream generator skips followups entirely when with_followups=False -""" -from __future__ import annotations - -import asyncio -import inspect -import json -import sys -import unittest.mock as _mock -from pathlib import Path - -_BACKEND_ROOT = Path(__file__).resolve().parent.parent -if str(_BACKEND_ROOT) not in sys.path: - sys.path.insert(0, str(_BACKEND_ROOT)) - - -# ========================================================================= -# Standalone runner -# ========================================================================= - -PASS = "\033[92mPASS\033[0m" -FAIL = "\033[91mFAIL\033[0m" -_failures: list = [] - - -def check(name: str, cond: bool, detail: str = ""): - status = PASS if cond else FAIL - suffix = f" — {detail}" if detail and not cond else "" - print(f" [{status}] {name}{suffix}") - if not cond: - _failures.append(name) - - -# ========================================================================= -# 1. Pydantic model -# ========================================================================= - -def test_chat_request_default_with_followups_is_true(): - """Default ON — the whole point of the feature is to suggest next steps.""" - from app.models.chat import ChatRequest - req = ChatRequest(message="hi") - check("ChatRequest: with_followups defaults to True", - req.with_followups is True) - - -def test_chat_request_explicit_with_followups_false(): - """User can opt out (e.g. for a streaming benchmark run).""" - from app.models.chat import ChatRequest - req = ChatRequest(message="hi", with_followups=False) - check("ChatRequest: with_followups=False round-trips", - req.with_followups is False) - - -def test_chat_response_has_followups_field(): - """Followups default to an empty list so the field is always present - in the response envelope (no 'undefined' in the frontend).""" - from app.models.chat import ChatResponse - resp = ChatResponse(message="x", conversation_id="c1") - check("ChatResponse: followups field exists", hasattr(resp, "followups")) - check("ChatResponse: followups defaults to []", - resp.followups == []) - - -# ========================================================================= -# 2. LLMService.generate_followups -# ========================================================================= - -def test_llm_service_has_generate_followups_method(): - from app.services.llm import LLMService - sig = inspect.signature(LLMService.generate_followups) - params = sig.parameters - check("LLMService.generate_followups exists", - hasattr(LLMService, "generate_followups")) - check("LLMService.generate_followups is async", - inspect.iscoroutinefunction(LLMService.generate_followups)) - check("LLMService.generate_followups has 'query' param", - "query" in params) - check("LLMService.generate_followups has 'answer' param", - "answer" in params) - check("LLMService.generate_followups has 'n' param with default 3", - "n" in params and params["n"].default == 3) - - -def test_generate_followups_parses_clean_json(): - """The LLM returns a clean JSON array — parse it as-is.""" - from app.services.llm import LLMService - svc = LLMService() - raw = json.dumps([ - "What is the difference between QLoRA and LoRA?", - "Which papers introduced retrieval-augmented generation?", - "How does the reranker work in this system?", - ]) - with _mock.patch.object(svc, "chat_complete", _mock.AsyncMock(return_value=raw)): - out = asyncio.run(svc.generate_followups( - query="Explain QLoRA", answer="QLoRA is …")) - check("generate_followups: returns 3 strings from clean JSON", - len(out) == 3 - and out[0].startswith("What is the difference")) - check("generate_followups: each entry is a non-empty string", - all(isinstance(s, str) and s.strip() for s in out)) - - -def test_generate_followups_parses_json_with_surrounding_prose(): - """LLMs often wrap JSON in ```json ... ``` fences or preambles — we - must extract the array even when the response isn't pure JSON.""" - from app.services.llm import LLMService - svc = LLMService() - raw = ( - "Here are 3 follow-up questions:\n" - "```json\n" - + json.dumps([ - "How does chunking work?", - "What is the embedding model?", - "Why use Neo4j?", - ]) - + "\n```\n" - ) - with _mock.patch.object(svc, "chat_complete", _mock.AsyncMock(return_value=raw)): - out = asyncio.run(svc.generate_followups(query="q", answer="a")) - check("generate_followups: extracts array from fenced prose", - len(out) == 3 and out[0] == "How does chunking work?") - - -def test_generate_followups_caps_at_n_and_drops_blanks(): - """If the LLM returns 5 (n=3) or has empty strings, we cap and clean.""" - from app.services.llm import LLMService - svc = LLMService() - raw = json.dumps([ - "Good question 1", - "", - " ", - "Good question 2", - "Good question 3", - ]) - with _mock.patch.object(svc, "chat_complete", _mock.AsyncMock(return_value=raw)): - out = asyncio.run(svc.generate_followups(query="q", answer="a", n=3)) - check("generate_followups: caps at n=3", len(out) == 3) - check("generate_followups: drops blank entries", - out == ["Good question 1", "Good question 2", "Good question 3"]) - - -def test_generate_followups_returns_empty_on_llm_error(): - """A failure here must NEVER propagate — followups are a UX bonus. - The chat response should still be valid (just with no chips).""" - from app.services.llm import LLMService - svc = LLMService() - with _mock.patch.object( - svc, "chat_complete", - _mock.AsyncMock(side_effect=RuntimeError("rate limit")), - ): - out = asyncio.run(svc.generate_followups(query="q", answer="a")) - check("generate_followups: returns [] on LLM error", out == []) - - -def test_generate_followups_returns_empty_on_garbage(): - """If the LLM returns text with no JSON array at all, we don't crash.""" - from app.services.llm import LLMService - svc = LLMService() - with _mock.patch.object( - svc, "chat_complete", - _mock.AsyncMock(return_value="I cannot generate followups right now."), - ): - out = asyncio.run(svc.generate_followups(query="q", answer="a")) - check("generate_followups: returns [] when no JSON in response", - out == []) - - -def test_generate_followups_handles_fewer_than_n(): - """LLM sometimes returns only 2 — return what we got (length 2), - the client renders what it gets (no padding).""" - from app.services.llm import LLMService - svc = LLMService() - raw = json.dumps(["Q1", "Q2"]) - with _mock.patch.object(svc, "chat_complete", _mock.AsyncMock(return_value=raw)): - out = asyncio.run(svc.generate_followups(query="q", answer="a", n=3)) - check("generate_followups: returns 2 when LLM gives 2", len(out) == 2) - - -# ========================================================================= -# 3. chat_stream_generator emits a 'followups' SSE event -# ========================================================================= - -class _FakeStreamLLM: - """Stub for LLMService used by chat_stream_generator tests.""" - def __init__(self, full_text="answer body", followups=None, - followups_raises=False): - self._text = full_text - self._followups = followups or [] - self._raises = followups_raises - self.chat_complete_calls = 0 - self.generate_followups_calls = 0 - - async def chat_complete_stream(self, messages): - for chunk in self._text.split(" "): - yield chunk + " " - - async def generate_followups(self, query, answer, n=3): - self.generate_followups_calls += 1 - if self._raises: - raise RuntimeError("boom") - return list(self._followups) - - -def _collect_sse_events(generator): - """Drain the async generator and return the joined body.""" - out = [] - async def _drain(): - async for chunk in generator: - out.append(chunk) - asyncio.run(_drain()) - return "".join(out) - - -def _make_stub_db(): - """No-op DB context manager that returns a cursor returning empty results. - - aiosqlite's `db.execute()` is sync and returns an `AsyncCursor` that is - BOTH awaitable (for `await db.execute(INSERT...)` in non-SELECT paths) - AND an async context manager (for `async with db.execute(SELECT...)`). - The fake must support both usage patterns from chat.py.""" - class _CursorCtx: - # Context manager protocol: `async with db.execute(...) as cursor:` - async def __aenter__(self): return self - async def __aexit__(self, *a): return False - # Awaitable protocol: `await db.execute(INSERT/UPDATE/DELETE...)` - # (the production code uses await for non-SELECT writes) - def __await__(self): - async def _coro(): - return self - return _coro().__await__() - async def fetchall(self): return [] - async def fetchone(self): return None - @property - def rowcount(self): return 0 - - class _DB: - async def __aenter__(self): return self - async def __aexit__(self, *a): return False - def execute(self, *a, **kw): return _CursorCtx() - async def commit(self): return None - @property - def rowcount(self): return 0 - - def _stub(): - return _DB() - return _stub - - -def test_stream_generator_emits_followups_event(): - """When with_followups=True, a 'event: followups' SSE block should - appear in the stream with the list of questions as JSON data. - - Also covers the disabled case (with_followups=False) in the same - test. Both cases must run in ONE event loop — calling `asyncio.run` - twice in this process trips an httpx telemetry shutdown that raises - "Event loop is closed" on the second loop. - """ - from app.api import chat as chat_mod - from app.models.chat import ChatRequest - - fake_llm_on = _FakeStreamLLM( - full_text="streaming answer", - followups=["Q1?", "Q2?", "Q3?"], - ) - fake_llm_off = _FakeStreamLLM(full_text="x", followups=["Q1", "Q2", "Q3"]) - - async def _run_both(): - # Each case uses its own get_llm_service; swap inside the loop. - async def _drain(req, llm): - async def _stub(): - return llm - with _mock.patch.object(chat_mod, "get_db", _make_stub_db()), \ - _mock.patch.object(chat_mod, "get_llm_service", _stub): - gen = chat_mod.chat_stream_generator(req, user_id=1) - out = [] - async for chunk in gen: - out.append(chunk) - return "".join(out) - - body_on = await _drain( - ChatRequest(message="hi", with_followups=True), fake_llm_on, - ) - body_off = await _drain( - ChatRequest(message="hi", with_followups=False), fake_llm_off, - ) - return body_on, body_off - - body_on, body_off = asyncio.run(_run_both()) - - # ---- Case 1: with_followups=True ---- - check("stream: emits 'event: followups' when on", - "event: followups" in body_on) - check("stream: followups data is valid JSON with 3 items", - '"followups"' in body_on and '"Q1?"' in body_on and '"Q3?"' in body_on) - check("stream: still emits 'event: done' when on", - "event: done" in body_on) - check("stream: called generate_followups exactly once when on", - fake_llm_on.generate_followups_calls == 1) - # ---- Case 2: with_followups=False ---- - check("stream: no 'event: followups' when with_followups=False", - "event: followups" not in body_off) - check("stream: generate_followups not called when off", - fake_llm_off.generate_followups_calls == 0) - - -# ========================================================================= -# Driver -# ========================================================================= - -ALL_TESTS = [ - test_chat_request_default_with_followups_is_true, - test_chat_request_explicit_with_followups_false, - test_chat_response_has_followups_field, - test_llm_service_has_generate_followups_method, - test_generate_followups_parses_clean_json, - test_generate_followups_parses_json_with_surrounding_prose, - test_generate_followups_caps_at_n_and_drops_blanks, - test_generate_followups_returns_empty_on_llm_error, - test_generate_followups_returns_empty_on_garbage, - test_generate_followups_handles_fewer_than_n, - test_stream_generator_emits_followups_event, -] - - -def main() -> int: - print(f"Running {len(ALL_TESTS)} checks for #0 follow-up suggestions...") - for fn in ALL_TESTS: - try: - fn() - except Exception as e: - check(f"{fn.__name__}: no unhandled exceptions", False, repr(e)) - print() - if _failures: - print(f"{FAIL} {len(_failures)} FAILED: " + ", ".join(_failures)) - return 1 - print(f"{PASS} All checks passed ({len(ALL_TESTS)} tests).") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/docker-compose.yml b/docker-compose.yml index ce4c63d..c50e7b0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,5 +1,3 @@ -version: '3.8' - services: neo4j: image: neo4j:5.14-community @@ -23,6 +21,25 @@ services: chromadb: image: chromadb/chroma:0.4.18 container_name: chromadb_kg + # The stock image entrypoint runs `pip install --force-reinstall + # chroma-hnswlib` on EVERY start — today's chroma-hnswlib pulls in + # numpy 2.x, which removed np.float_ and crashes chromadb 0.4.18 at + # import. Bypass the entrypoint and run uvicorn directly with the + # image's own (working) hnswlib build; same args as the original. + entrypoint: ["/usr/local/bin/uvicorn"] + command: + - "chromadb.app:app" + - "--workers" + - "1" + - "--host" + - "0.0.0.0" + - "--port" + - "8000" + - "--proxy-headers" + - "--log-config" + - "chromadb/log_config.yml" + - "--timeout-keep-alive" + - "30" ports: - "8000:8000" environment: diff --git a/frontend/src/api/chat.js b/frontend/src/api/chat.js index 2f0b20d..02c8f8f 100644 --- a/frontend/src/api/chat.js +++ b/frontend/src/api/chat.js @@ -1,18 +1,17 @@ import service from './index' export const chatApi = { - send: (message, conversationId, includeContext = true, useGraphRag = false, compareMode = false, withFollowups = true) => { + send: (message, conversationId, includeContext = true, useGraphRag = false, compareMode = false) => { return service.post('/chat', { message, conversation_id: conversationId, include_context: includeContext, use_graph_rag: useGraphRag, - compare_mode: compareMode, - with_followups: withFollowups + compare_mode: compareMode }) }, - stream: (message, conversationId, useGraphRag = false, compareMode = false) => { + stream: (message, conversationId, useGraphRag = false, compareMode = false, enableThinking = false) => { const token = localStorage.getItem('token') return fetch('/api/chat/stream', { method: 'POST', @@ -25,7 +24,8 @@ export const chatApi = { conversation_id: conversationId, include_context: true, use_graph_rag: useGraphRag, - compare_mode: compareMode + compare_mode: compareMode, + enable_thinking: enableThinking }) }) }, diff --git a/frontend/src/utils/sse.js b/frontend/src/utils/sse.js new file mode 100644 index 0000000..b5c0071 --- /dev/null +++ b/frontend/src/utils/sse.js @@ -0,0 +1,62 @@ +// Minimal Server-Sent Events parser for the streaming chat endpoint. +// +// Pure, framework-free helpers so the framing logic can be unit-tested in +// plain Node (see tests/test_sse.cjs, which mirrors these line-for-line). +// +// The backend (app/api/chat.py) emits standard SSE: named events via +// `event: ` plus one or more `data: ` lines, blocks separated by +// a blank line (`\n\n`). The default (unnamed) event carries body tokens as +// `data: {"chunk": "..."}`. + +/** + * Parse a single SSE block (the text between blank-line separators). + * Returns `{ event, data }` where `event` defaults to "message" and `data` + * is the JSON-parsed payload (or null if the block has no data / invalid + * JSON). Returns null for empty/heartbeat blocks. + */ +export function parseSseBlock(block) { + let eventName = 'message' + const dataLines = [] + for (const line of block.split('\n')) { + if (line.startsWith('event:')) eventName = line.slice(6).trim() + else if (line.startsWith('data:')) dataLines.push(line.slice(5).trim()) + } + if (dataLines.length === 0) return null + let payload = null + try { + payload = JSON.parse(dataLines.join('\n')) + } catch { + payload = null + } + return { event: eventName, data: payload } +} + +/** + * Create a streaming parser that buffers arbitrary text chunks (which may + * split SSE blocks mid-frame) and invokes `onEvent(event, data)` for each + * complete block. Call `flush()` once the stream ends to drain any trailing + * block that lacked a final blank line. + */ +export function createSseParser(onEvent) { + let buffer = '' + + const emit = (block) => { + const parsed = parseSseBlock(block) + if (parsed) onEvent(parsed.event, parsed.data) + } + + return { + feed(text) { + buffer += text + let sep + while ((sep = buffer.indexOf('\n\n')) !== -1) { + emit(buffer.slice(0, sep)) + buffer = buffer.slice(sep + 2) + } + }, + flush() { + if (buffer.trim()) emit(buffer) + buffer = '' + }, + } +} diff --git a/frontend/src/views/ChatPage.vue b/frontend/src/views/ChatPage.vue index 2e3e335..74dc316 100644 --- a/frontend/src/views/ChatPage.vue +++ b/frontend/src/views/ChatPage.vue @@ -2,8 +2,8 @@