fix: 分批写入 Milvus chunks 避免 gRPC 超限

backend/package/yuxi/knowledge/implementations/milvus.py

@@ -33,6 +33,7 @@ MILVUS_AVAILABLE = True
 CONTENT_SPARSE_FIELD = "content_sparse"
 CONTENT_ANALYZER_PARAMS = {"type": "chinese"}
 VECTOR_METRIC_TYPE = "COSINE"
+MILVUS_CHUNK_EMBED_BATCH_SIZE = 200
 
 
 @dataclass(kw_only=True)
@@ -526,6 +527,33 @@ class MilvusKB(KnowledgeBase):
             logger.error(f"Failed to rollback Milvus chunks for {file_id}: {cleanup_error}")
         raise errors[0]
 
+    async def _embed_and_store_chunks(
+        self,
+        kb_id: str,
+        file_id: str,
+        collection: Collection,
+        chunks: list[dict],
+        embedding_function,
+        *,
+        chunk_batch_size: int = MILVUS_CHUNK_EMBED_BATCH_SIZE,
+    ) -> None:
+        """对 chunks 进行分批嵌入并存储到 Milvus 和 PostgreSQL"""
+        if not chunks:
+            return
+
+        chunk_batch_size = max(int(chunk_batch_size), 1)
+        for start in range(0, len(chunks), chunk_batch_size):
+            batch_chunks = chunks[start : start + chunk_batch_size]
+            texts = [chunk["content"] for chunk in batch_chunks]
+            embeddings = await embedding_function(texts)
+            await self._insert_chunks_to_stores(
+                kb_id,
+                file_id,
+                collection,
+                batch_chunks,
+                embeddings,
+            )
+
     async def _delete_file_chunks_from_milvus(self, collection: Collection, file_id: str) -> None:
         expr = f'file_id == "{file_id}"'
         results = collection.query(expr=expr, output_fields=["id"], limit=1)
@@ -652,12 +680,9 @@ class MilvusKB(KnowledgeBase):
             )
 
             if chunks:
-                texts = [chunk["content"] for chunk in chunks]
-                embeddings = await embedding_function(texts)
-
                 # Clean up existing chunks if any (for re-indexing)
                 await self.delete_file_chunks_only(kb_id, file_id)
-                await self._insert_chunks_to_stores(kb_id, file_id, collection, chunks, embeddings)
+                await self._embed_and_store_chunks(kb_id, file_id, collection, chunks, embedding_function)
 
             logger.info(f"Indexed file {file_id} into Milvus")
 
@@ -744,9 +769,7 @@ class MilvusKB(KnowledgeBase):
                 await self.delete_file_chunks_only(kb_id, file_id)
 
                 if chunks:
-                    texts = [chunk["content"] for chunk in chunks]
-                    embeddings = await embedding_function(texts)
-                    await self._insert_chunks_to_stores(kb_id, file_id, collection, chunks, embeddings)
+                    await self._embed_and_store_chunks(kb_id, file_id, collection, chunks, embedding_function)
 
                 logger.info(f"Updated file {file_path} in Milvus. Done.")
 

backend/test/unit/plugins/test_milvus_kb.py

@@ -1,5 +1,14 @@
+import asyncio
+import types
+
+import pytest
 from pymilvus import CollectionSchema, DataType, FieldSchema, Function, FunctionType
 
+import yuxi
+
+if "knowledge_base" not in vars(yuxi):
+    yuxi.knowledge_base = types.SimpleNamespace()
+
 from yuxi.knowledge.implementations.milvus import (
     CONTENT_ANALYZER_PARAMS,
     CONTENT_SPARSE_FIELD,
@@ -23,6 +32,7 @@ class FakeCollection:
     def __init__(self, distance: float = 0.8):
         self.search_calls = []
         self.hybrid_calls = []
+        self.insert_calls = []
         self.distance = distance
 
     def search(self, **kwargs):
@@ -33,6 +43,9 @@ class FakeCollection:
         self.hybrid_calls.append(kwargs)
         return [[FakeHit("Hybrid result", self.distance)]]
 
+    def insert(self, entities):
+        self.insert_calls.append(entities)
+
 
 def make_kb(collection: FakeCollection) -> MilvusKB:
     kb = MilvusKB.__new__(MilvusKB)
@@ -48,6 +61,16 @@ def make_kb(collection: FakeCollection) -> MilvusKB:
     return kb
 
 
+def make_chunk(index: int, content: str = "content") -> dict:
+    return {
+        "id": f"id-{index}",
+        "chunk_id": f"chunk-{index}",
+        "file_id": "file-1",
+        "chunk_index": index,
+        "content": content,
+    }
+
+
 def test_build_chunk_pg_records_preserves_extraction_result():
     kb = MilvusKB.__new__(MilvusKB)
 
@@ -67,6 +90,179 @@ def test_build_chunk_pg_records_preserves_extraction_result():
     assert records[0]["extraction_result"] == {"entities": ["alpha"]}
 
 
+async def test_embed_and_store_chunks_batches_embedding_and_insert():
+    kb = MilvusKB.__new__(MilvusKB)
+    chunks = [make_chunk(index, content=f"text-{index}") for index in range(450)]
+    embedding_calls = []
+    store_calls = []
+
+    async def embedding_function(texts):
+        embedding_calls.append(list(texts))
+        return [[float(len(text))] for text in texts]
+
+    async def insert_chunks_to_stores(kb_id, file_id, collection, batch_chunks, embeddings, **kwargs):
+        store_calls.append(
+            {
+                "kb_id": kb_id,
+                "file_id": file_id,
+                "chunks": list(batch_chunks),
+                "embeddings": list(embeddings),
+                "kwargs": kwargs,
+            }
+        )
+
+    kb._insert_chunks_to_stores = insert_chunks_to_stores
+
+    await kb._embed_and_store_chunks(
+        "db",
+        "file-1",
+        FakeCollection(),
+        chunks,
+        embedding_function,
+        chunk_batch_size=200,
+    )
+
+    assert [len(call) for call in embedding_calls] == [200, 200, 50]
+    assert [len(call["chunks"]) for call in store_calls] == [200, 200, 50]
+    assert store_calls[0]["chunks"][0]["chunk_id"] == "chunk-0"
+    assert store_calls[1]["chunks"][0]["chunk_id"] == "chunk-200"
+    assert store_calls[2]["chunks"][0]["chunk_id"] == "chunk-400"
+    assert all(call["kwargs"] == {} for call in store_calls)
+
+
+async def test_insert_chunks_to_stores_inserts_current_batch(monkeypatch):
+    repos = []
+
+    class FakeChunkRepo:
+        def __init__(self):
+            self.upsert_calls = []
+            self.delete_calls = []
+            repos.append(self)
+
+        async def batch_upsert(self, chunks):
+            self.upsert_calls.append(chunks)
+            return []
+
+        async def delete_by_file_id(self, file_id):
+            self.delete_calls.append(file_id)
+            return 0
+
+    monkeypatch.setattr("yuxi.knowledge.implementations.milvus.KnowledgeChunkRepository", FakeChunkRepo)
+    kb = MilvusKB.__new__(MilvusKB)
+    collection = FakeCollection()
+    chunks = [make_chunk(index) for index in range(3)]
+    embeddings = [[0.1, 0.2] for _ in chunks]
+
+    await kb._insert_chunks_to_stores("db", "file-1", collection, chunks, embeddings)
+
+    assert len(collection.insert_calls) == 1
+    assert collection.insert_calls[0][0] == ["id-0", "id-1", "id-2"]
+    assert collection.insert_calls[0][5] == embeddings
+    assert len(repos[0].upsert_calls) == 1
+    assert [record["chunk_id"] for record in repos[0].upsert_calls[0]] == ["chunk-0", "chunk-1", "chunk-2"]
+
+
+async def test_insert_chunks_to_stores_rolls_back_file_when_milvus_insert_fails(monkeypatch):
+    repos = []
+
+    class FakeChunkRepo:
+        def __init__(self):
+            self.upsert_calls = []
+            self.delete_calls = []
+            repos.append(self)
+
+        async def batch_upsert(self, chunks):
+            self.upsert_calls.append(chunks)
+            return []
+
+        async def delete_by_file_id(self, file_id):
+            self.delete_calls.append(file_id)
+            return 0
+
+    class FailingCollection(FakeCollection):
+        def insert(self, entities):
+            super().insert(entities)
+            raise RuntimeError("milvus boom")
+
+    monkeypatch.setattr("yuxi.knowledge.implementations.milvus.KnowledgeChunkRepository", FakeChunkRepo)
+    kb = MilvusKB.__new__(MilvusKB)
+    collection = FailingCollection()
+    milvus_delete_calls = []
+
+    async def delete_file_chunks_from_milvus(collection_arg, file_id):
+        milvus_delete_calls.append((collection_arg, file_id))
+
+    kb._delete_file_chunks_from_milvus = delete_file_chunks_from_milvus
+    chunks = [make_chunk(index) for index in range(2)]
+    embeddings = [[0.1, 0.2] for _ in chunks]
+
+    with pytest.raises(RuntimeError, match="milvus boom"):
+        await kb._insert_chunks_to_stores("db", "file-1", collection, chunks, embeddings)
+
+    assert repos[0].delete_calls == ["file-1"]
+    assert milvus_delete_calls == [(collection, "file-1")]
+
+
+async def test_update_content_uses_streaming_chunk_store(monkeypatch):
+    kb = MilvusKB.__new__(MilvusKB)
+    kb.databases_meta = {"db": {"embedding_model_spec": "test-provider:test-embedding", "metadata": {}}}
+    kb.files_meta = {
+        "file-1": {
+            "path": "/tmp/demo.md",
+            "filename": "demo.md",
+            "processing_params": {},
+        }
+    }
+    kb._metadata_lock = asyncio.Lock()
+    collection = FakeCollection()
+    queue_adds = []
+    queue_removes = []
+    persisted_files = []
+    deleted_files = []
+    store_calls = []
+
+    async def get_collection(kb_id):
+        return collection
+
+    async def forbidden_embedding(texts):
+        raise AssertionError("update_content should not embed the whole file directly")
+
+    async def persist_file(file_id):
+        persisted_files.append(file_id)
+
+    async def delete_file_chunks_only(kb_id, file_id):
+        deleted_files.append((kb_id, file_id))
+
+    async def embed_and_store_chunks(kb_id, file_id, collection_arg, chunks, embedding_function):
+        store_calls.append((kb_id, file_id, collection_arg, list(chunks), embedding_function))
+
+    async def parse_file(source, params):
+        return "# markdown"
+
+    kb._get_milvus_collection = get_collection
+    kb._get_embedding_function = lambda embedding_model_spec: forbidden_embedding
+    kb._persist_file = persist_file
+    kb._split_text_into_chunks = lambda text, file_id, filename, params: [make_chunk(0), make_chunk(1)]
+    kb.delete_file_chunks_only = delete_file_chunks_only
+    kb._embed_and_store_chunks = embed_and_store_chunks
+    kb._add_to_processing_queue = lambda file_id: queue_adds.append(file_id)
+    kb._remove_from_processing_queue = lambda file_id: queue_removes.append(file_id)
+    monkeypatch.setattr("yuxi.knowledge.implementations.milvus.Parser.aparse", parse_file)
+
+    result = await kb.update_content("db", ["file-1"])
+
+    assert deleted_files == [("db", "file-1")]
+    assert len(store_calls) == 1
+    assert store_calls[0][2] is collection
+    assert [chunk["chunk_id"] for chunk in store_calls[0][3]] == ["chunk-0", "chunk-1"]
+    assert store_calls[0][4] is forbidden_embedding
+    assert result[0]["status"] == "done"
+    assert kb.files_meta["file-1"]["status"] == "done"
+    assert queue_adds == ["file-1"]
+    assert queue_removes == ["file-1"]
+    assert persisted_files
+
+
 async def test_keyword_mode_uses_milvus_bm25_search():
     collection = FakeCollection()
     kb = make_kb(collection)

docs/develop-guides/changelog.md

@@ -29,7 +29,7 @@
 - 重构智能体配置语义：用户可见的 `AgentConfig` 收敛为数据库持久化的一级 `Agent`，内置 Python Agent 改为智能体后端；新增 `/api/agent` 管理与运行接口，聊天、运行任务、恢复审批和文件预览均从线程绑定的 Agent 解析运行时上下文，前端只提交 `agent_id`，并在模型配置页新增“智能体”管理页签。
 - 删除 Upload 与 LightRAG 图谱/知识库能力：知识库类型收敛为 Milvus 与 Dify，只保留 Milvus 知识库内图谱构建/展示/检索，移除独立 `/graph` 页面和默认上传图谱工具。
 - 收敛只读知识源连接器：新增 `ReadOnlyConnectors` 基类，Dify 改为声明自身创建参数与校验规则，新增 Notion Data Source 只读知识库并支持 Search/Find/Open；知识库类型接口返回创建参数 schema，前端新建表单按类型动态渲染非 Milvus 配置并统一保存到 `additional_params`。
-- 新增知识库 Chunk 持久化：Milvus 知识库索引/更新流程会将 chunks 双写到 PostgreSQL `knowledge_chunks` 表与 Milvus，文件内容查看优先查询 PostgreSQL，并为位置信息、图谱实体关联、标签和抽取结果预留结构化字段。
+- 新增知识库 Chunk 持久化：Milvus 知识库索引/更新流程会将 chunks 双写到 PostgreSQL `knowledge_chunks` 表与 Milvus，文件内容查看优先查询 PostgreSQL，并为位置信息、图谱实体关联、标签和抽取结果预留结构化字段；chunk 入库改为分批 embedding 与分批写入，避免大文件一次性写入触发 gRPC 消息大小限制。
 - 完善 Milvus 知识库图谱构建：修复 Chunk 图谱写入返回值、Neo4j 同步写入阻塞事件循环、重复构建任务竞态、图谱查询提前终止、Neo4j 连接复用、LLM 抽取超时重试和前端错误详情展示等问题；图谱构建会将 entity/triple 本体与 chunk 引用写入 PostgreSQL，并为唯一 entity/triple 建立 Milvus 语义索引，单文件删除时同步清理图谱引用和孤儿向量。
 - 优化图谱抽取器配置：未配置时在图谱中心展示配置入口，抽取方案收敛为 LLM，前端仅保留“更多拓展中”占位；LLM 抽取器使用固定 Prompt + 自定义 Schema，并支持模型参数与并发队列数；已配置后允许修改参数并提示重置重抽风险。修复上传并入库新文件时旧内存 metadata 覆盖数据库图谱配置的问题。
 - 新增 Milvus 图谱检索链路：Query 可召回图谱实体和三元组，结合 Chunk 命中实体构造 seed entity，读取 Neo4j 2-hop 子图后用 igraph 执行 PPR，最终以 Chunk 为产物并通过 RRF 与原 Chunk 召回融合；检索配置改为 dataclass 元数据生成，支持 `depend_on` 控制重排序和图检索参数展示。