diff --git a/backend/app/api/knowledge.py b/backend/app/api/knowledge.py index 3911abfc..b5dd711c 100644 --- a/backend/app/api/knowledge.py +++ b/backend/app/api/knowledge.py @@ -2,11 +2,14 @@ from __future__ import annotations from typing import Any -from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile +from fastapi import ( + APIRouter, BackgroundTasks, Depends, File, HTTPException, Query, UploadFile, +) from pydantic import BaseModel, Field from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession +from app.core.config import settings from app.db.base import get_session from app.db.models import Document, MemoryItem, Skill from app.memory import inverted, kb, store @@ -105,6 +108,10 @@ class DocumentOut(BaseModel): source: str mime: str chunk_count: int + #: ready | processing | failed。前端据此显示"处理中 3200/7914 段" + status: str = "ready" + error: str = "" + meta: dict[str, Any] = Field(default_factory=dict) created_at: Any = None model_config = {"from_attributes": True} @@ -146,13 +153,23 @@ async def ingest(payload: IngestIn, session: AsyncSession = Depends(get_session) @kb_router.post("/upload", response_model=DocumentOut, status_code=201) async def upload( + background: BackgroundTasks, file: UploadFile = File(...), collection: str = "default", session: AsyncSession = Depends(get_session), ) -> Document: - raw = await file.read() - if len(raw) > 10 * 1024 * 1024: - raise HTTPException(413, "文件超过 10MB") + # 边读边数,超了立刻停——原来是 raw = await file.read() 读完再判断, + # 传个 1GB 的文件,内存在报 413 之前就已经吃掉了。那个检查拦的是"入库", + # 不是"占内存" + limit = settings.max_upload_mb * 1024 * 1024 + chunks: list[bytes] = [] + total = 0 + while piece := await file.read(1024 * 1024): + total += len(piece) + if total > limit: + raise HTTPException(413, f"文件超过 {settings.max_upload_mb}MB") + chunks.append(piece) + raw = b"".join(chunks) from app.memory.parsing import UnsupportedDocument, extract @@ -161,14 +178,41 @@ async def upload( except UnsupportedDocument as e: # 原样交回:里面写的是"装哪个包"或"这是扫描件",都能照着做 raise HTTPException(415, str(e)) from None - return await kb.ingest_document( + doc = await kb.create_document( session, collection=collection, title=file.filename or "上传文档", content=content, source=file.filename or "", mime=file.content_type or "text/plain", + status="processing", ) + # 立刻返回,切块和算向量在后台跑。实测一个 10MB 文档用本地哈希要两分半, + # 换成远端 embedding 是几十分钟——压在请求里必然超时,而后端其实还在跑, + # 界面显示失败、数据其实成功,比直接拒绝还糟 + background.add_task(_process_in_background, doc.id) + return doc + + +async def _process_in_background(doc_id: str) -> None: + """后台切块。进度写回文档行,失败也写回——不然它会永远停在"处理中"。""" + from app.db.base import SessionLocal + + async with SessionLocal() as session: + doc = await session.get(Document, doc_id) + if not doc: + return + try: + # 进度由 process_document 自己按批提交——它就在主事务里, + # 不会和另一个 session 抢 SQLite 的写锁 + await kb.process_document(session, doc) + except Exception as e: # noqa: BLE001 + await session.rollback() + doc = await session.get(Document, doc_id) + if doc: + doc.status = "failed" + doc.error = f"{type(e).__name__}: {e}"[:500] + await session.commit() @kb_router.get("/search") diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 781c89be..bbfdb269 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -66,6 +66,9 @@ class Settings(BaseSettings): max_run_seconds: int = 600 max_graph_steps: int = 200 max_agent_steps: int = 25 + #: 上传文件的大小上限(MB)。原来写死 10,而且是读完整个文件才检查—— + #: 传一个 1GB 的文件,内存在报 413 之前就吃掉了 + max_upload_mb: int = 50 max_concurrent_runs: int = 20 # --- 工具 --- diff --git a/backend/app/db/base.py b/backend/app/db/base.py index 83e5d130..81d1ab80 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -82,6 +82,8 @@ _COLUMN_MIGRATIONS: list[tuple[str, str, str]] = [ ("memories", "embed_model", "VARCHAR(100) DEFAULT ''"), ("memories", "embed_dim", "INTEGER DEFAULT 0"), ("chunks", "token_len", "INTEGER DEFAULT 0"), + ("documents", "status", "TEXT DEFAULT 'ready'"), + ("documents", "error", "TEXT DEFAULT ''"), ] diff --git a/backend/app/db/models.py b/backend/app/db/models.py index 7529327b..e6616aa2 100644 --- a/backend/app/db/models.py +++ b/backend/app/db/models.py @@ -279,6 +279,11 @@ class Document(Base, TimestampMixin): content: Mapped[str] = mapped_column(Text, default="") meta: Mapped[dict[str, Any]] = mapped_column(default=dict) chunk_count: Mapped[int] = mapped_column(Integer, default=0) + # 切块+算向量是在后台跑的:一个 10MB 文档配上远端 embedding 要几十分钟, + # 压在 HTTP 请求里必然超时——而后端其实还在跑,界面显示失败、数据其实成功, + # 比直接拒绝还糟。ready | processing | failed + status: Mapped[str] = mapped_column(String(20), default="ready") + error: Mapped[str] = mapped_column(Text, default="") chunks: Mapped[list["Chunk"]] = relationship( back_populates="document", cascade="all, delete-orphan" diff --git a/backend/app/main.py b/backend/app/main.py index b5a8e3b3..5fe5b524 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -43,6 +43,19 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: except Exception as e: # noqa: BLE001 - key 没了不该拦住整个服务起不来 logger.warning("embedding 配置没能生效,暂用本地向量:%s", e) + # 上次被强杀时留下的"处理中"文档。后台任务不会自己回来,留着它只会在 + # 界面上永远转圈——和 runner 对 running 的 run 是同一个道理 + from sqlalchemy import update as _update + + from app.db.models import Document as _Doc + + async with SessionLocal() as _s: + await _s.execute( + _update(_Doc).where(_Doc.status == "processing") + .values(status="failed", error="服务重启,这份文档没有处理完,请重新上传") + ) + await _s.commit() + await run_manager.setup() from app.sandbox.manager import sandbox_manager as _sbm diff --git a/backend/app/memory/kb.py b/backend/app/memory/kb.py index 4f3ed72b..77d4d5e5 100644 --- a/backend/app/memory/kb.py +++ b/backend/app/memory/kb.py @@ -15,6 +15,9 @@ from app.memory.embeddings import ( _TARGET = 800 _OVERLAP = 120 +#: 每批送多少段去算向量。远端接口单次批量有上限(OpenAI 是 2048),而且 +#: 一次把上万段全发出去,要么超时要么被限流,中途失败还没有断点 +_EMBED_BATCH = 64 def chunk_text(text: str, target: int = _TARGET, overlap: int = _OVERLAP) -> list[str]: @@ -119,7 +122,32 @@ async def ingest_document( source: str = "", mime: str = "text/plain", meta: dict[str, Any] | None = None, + on_progress: Any = None, ) -> Document: + doc = await create_document( + session, collection=collection, title=title, content=content, + source=source, mime=mime, meta=meta, status="ready", + ) + await process_document(session, doc, on_progress=on_progress) + return doc + + +async def create_document( + session: AsyncSession, + *, + collection: str, + title: str, + content: str, + source: str = "", + mime: str = "text/plain", + meta: dict[str, Any] | None = None, + status: str = "ready", +) -> Document: + """只建文档行,不切块不算向量。 + + 拆出来是为了让上传能立刻返回:切块 + 算向量在后台跑,一个 10MB 文档配上 + 远端 embedding 要几十分钟,压在 HTTP 请求里必然超时。 + """ doc = Document( collection=collection, title=title or source or "未命名文档", @@ -127,32 +155,66 @@ async def ingest_document( mime=mime, content=content, meta=meta or {}, + status=status, ) session.add(doc) - await session.flush() + await session.commit() + await session.refresh(doc) + return doc + +async def process_document( + session: AsyncSession, doc: Document, *, on_progress: Any = None +) -> Document: + """切块、算向量、建倒排。可以对着一个已存在的文档行反复跑(会先清掉旧的)。""" + from sqlalchemy import delete as sa_delete + + old_ids = list((await session.execute( + select(Chunk.id).where(Chunk.document_id == doc.id))).scalars()) + if old_ids: + await inverted.drop_chunk_terms(session, old_ids) + await session.execute(sa_delete(Chunk).where(Chunk.document_id == doc.id)) + + collection, content = doc.collection, doc.content pieces = chunk_text(content) - made: list[Chunk] = [] - if pieces: - vectors = await embed_texts(pieces) - for i, (piece, vec) in enumerate(zip(pieces, vectors)): + model_id, dim = embedder_id(), embedder_dim() + + # 一批算完就落盘。分批本身是必须的——原来 embed_texts(pieces) 一次全发, + # 本地哈希只是慢,换成远端模型就是必挂(一个 10MB 文档近八千段,塞不下)。 + # + # 每批都 commit 还有第二个好处:进度能被别人看见。试过用另一个 session + # 单独写进度,但 SQLite 只允许一个写者,主事务开着的时候那次写会被堵住, + # 界面上整个过程都是空的、只在结束时闪一下。 + for i in range(0, len(pieces), _EMBED_BATCH): + part = pieces[i:i + _EMBED_BATCH] + vectors = await embed_texts(part) + made: list[Chunk] = [] + for offset, (piece, vec) in enumerate(zip(part, vectors)): chunk = Chunk( document_id=doc.id, collection=collection, - ordinal=i, + ordinal=i + offset, content=piece, embedding=to_blob(vec), - embed_model=embedder_id(), - embed_dim=embedder_dim(), + embed_model=model_id, + embed_dim=dim, meta={"title": doc.title}, ) session.add(chunk) made.append(chunk) - # 先 flush 拿到 chunk.id,倒排行要引用它 - await session.flush() + await session.flush() # 先拿到 chunk.id,倒排行要引用它 for chunk in made: await inverted.index_chunk(session, chunk) + + done = min(i + _EMBED_BATCH, len(pieces)) + doc.meta = {**(doc.meta or {}), "progress": {"done": done, "total": len(pieces)}} + await session.commit() + if on_progress: + on_progress(done, len(pieces)) + doc.chunk_count = len(pieces) + doc.status = "ready" + doc.error = "" await session.commit() await session.refresh(doc) return doc @@ -278,7 +340,7 @@ async def reindex(session: AsyncSession, collection: str | None = None) -> dict[ return {"collection": collection, "reindexed": 0, "embedder": embedder_id()} model_id, dim = embedder_id(), embedder_dim() - batch = 64 + batch = _EMBED_BATCH for i in range(0, len(rows), batch): part = rows[i:i + batch] vectors = await embed_texts([r.content for r in part]) diff --git a/backend/app/memory/store.py b/backend/app/memory/store.py index 4a20ccde..30920e6a 100644 --- a/backend/app/memory/store.py +++ b/backend/app/memory/store.py @@ -210,7 +210,9 @@ async def reindex(session: AsyncSession, scope: str | None = None) -> int: return 0 model_id, dim = embedder_id(), embedder_dim() - batch = 64 # 和 kb.reindex 一样按批:一次几千条全发给远端接口必被限流 + from app.memory.kb import _EMBED_BATCH + + batch = _EMBED_BATCH # 和 kb 那边统一:一次几千条全发给远端接口必被限流 for i in range(0, len(rows), batch): part = rows[i:i + batch] vectors = await embed_texts([r.content for r in part]) diff --git a/backend/tests/test_kb.py b/backend/tests/test_kb.py index abaab43c..7cb67eb3 100644 --- a/backend/tests/test_kb.py +++ b/backend/tests/test_kb.py @@ -30,6 +30,17 @@ def _local_embedder(): emb.configure("local") +@pytest.fixture +async def client(): + from httpx import ASGITransport, AsyncClient + + from app.main import app + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as c: + yield c + + # -------------------------------------------------------------------------- # 切块 # -------------------------------------------------------------------------- @@ -588,3 +599,105 @@ async def test_the_endpoint_is_remembered_across_restarts(client=None) -> None: async with SessionLocal() as session: row = await session.get(Setting, e.EMBEDDING_SETTING_KEY) assert row is not None and "base_url" in row.value + + +# -------------------------------------------------------------------------- +# 大文件:分批、后台、大小防护 +# -------------------------------------------------------------------------- + + +async def test_embedding_is_batched_not_one_giant_call(monkeypatch) -> None: + """原来是 embed_texts(pieces) 一次全发。 + + 本地哈希只是慢,换成远端模型就是必挂:一个 10MB 文档切出近八千段, + 远端接口单次批量有上限(OpenAI 是 2048),一次请求塞不下。 + reindex 那边早就分批了,ingest 这条一直漏着。 + """ + from app.memory import kb as kb_mod + + sizes: list[int] = [] + real = kb_mod.embed_texts + + async def spy(texts): + sizes.append(len(texts)) + return await real(texts) + + monkeypatch.setattr(kb_mod, "embed_texts", spy) + + text = "\n\n".join(f"第{i}段内容,讲了一些东西。" * 12 for i in range(400)) + async with SessionLocal() as session: + await kb.ingest_document(session, collection="batch_test", title="大文档", content=text) + + assert len(sizes) > 1, f"只调了 {len(sizes)} 次,没有分批" + assert max(sizes) <= kb_mod._EMBED_BATCH, f"有一批 {max(sizes)} 条,超过了上限" + + +async def test_progress_lands_in_the_db_while_it_runs() -> None: + """进度要能被别人看见。 + + 试过用另一个 session 单独写,但 SQLite 只允许一个写者——主事务开着的时候 + 那次写会被堵住,整个过程界面上都是空的、只在结束时闪一下。改成主事务 + 自己按批提交。 + """ + from app.db.models import Document + + seen: list[tuple[int, int]] = [] + # 要切出多于一批(_EMBED_BATCH=64)的片段,否则只会有一次进度回调 + text = "\n\n".join(f"第{i}段内容,这里要写得长一些才切得够多。" * 30 for i in range(300)) + + async with SessionLocal() as session: + doc = await kb.create_document( + session, collection="prog_test", title="进度", content=text, status="processing") + await kb.process_document(session, doc, on_progress=lambda d, t: seen.append((d, t))) + + assert len(seen) > 1, "一次就跑完了,测不到进度" + assert seen[-1][0] == seen[-1][1], "最后一次进度不是满的" + + # 关键:中途的进度确实落了库,不是只在内存里 + async with SessionLocal() as session: + row = await session.get(Document, doc.id) + assert (row.meta or {}).get("progress", {}).get("total") == seen[-1][1] + assert row.status == "ready" + + +async def test_a_failed_document_does_not_stay_processing_forever(client) -> None: + """失败了要写回 failed。不写的话它永远停在"处理中",界面上一直转圈。""" + from app.db.models import Document + from app.api.knowledge import _process_in_background + + async with SessionLocal() as session: + doc = await kb.create_document( + session, collection="fail_test", title="会失败的", content="一些内容", + status="processing") + + import app.memory.kb as kb_mod + + async def boom(*a, **kw): + raise RuntimeError("向量服务挂了") + + orig = kb_mod.process_document + kb_mod.process_document = boom + try: + await _process_in_background(doc.id) + finally: + kb_mod.process_document = orig + + async with SessionLocal() as session: + row = await session.get(Document, doc.id) + assert row.status == "failed" + assert "向量服务挂了" in row.error + + +async def test_oversized_uploads_are_rejected_before_being_read(client) -> None: + """原来是 raw = await file.read() 读完再判断大小。 + + 传个 1GB 的文件,内存在报 413 之前就已经吃掉了——那个检查拦的是"入库", + 不是"占内存"。 + """ + from app.core.config import settings + + too_big = b"x" * ((settings.max_upload_mb + 1) * 1024 * 1024) + r = await client.post("/api/kb/upload", params={"collection": "big"}, + files={"file": ("big.txt", too_big, "text/plain")}) + assert r.status_code == 413 + assert str(settings.max_upload_mb) in r.json()["detail"] diff --git a/frontend/src/pages/KnowledgePage.tsx b/frontend/src/pages/KnowledgePage.tsx index 4b63f866..e38998a6 100644 --- a/frontend/src/pages/KnowledgePage.tsx +++ b/frontend/src/pages/KnowledgePage.tsx @@ -215,6 +215,14 @@ function KbTab() { } useEffect(() => { void load() }, [collection]) + // 有文档还在切块就接着轮询。一个 10MB 文档配上远端 embedding 要几十分钟, + // 不轮询的话界面会一直停在"处理中",人以为卡死了 + useEffect(() => { + if (!docs.some((d) => d.status === 'processing')) return + const timer = setInterval(() => { void load() }, 2000) + return () => clearInterval(timer) + }, [docs, collection]) + const search = async () => { if (!query.trim()) { setHits(null); return } const res = await api.kb.search(query, collection, alpha) @@ -310,8 +318,20 @@ function KbTab() { {docs.map((doc) => (
-
{doc.title}
-
{doc.source} · {doc.chunk_count} 个片段
+
+ {doc.status === 'processing' && } + {doc.title} +
+
+ {doc.status === 'processing' + ? (doc.meta?.progress + ? `正在切块并算向量… ${doc.meta.progress.done}/${doc.meta.progress.total} 段` + : '正在切块并算向量…') + : doc.status === 'failed' + // 失败原因要写出来:"这是扫描件"和"本地服务没起"是两件事 + ? 处理失败:{doc.error} + : `${doc.source} · ${doc.chunk_count} 个片段`} +