diff --git a/backend/package/yuxi/knowledge/implementations/milvus.py b/backend/package/yuxi/knowledge/implementations/milvus.py index b82b8bc1..2c89901e 100644 --- a/backend/package/yuxi/knowledge/implementations/milvus.py +++ b/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.") diff --git a/backend/test/unit/plugins/test_milvus_kb.py b/backend/test/unit/plugins/test_milvus_kb.py index 67be37aa..8d583eb1 100644 --- a/backend/test/unit/plugins/test_milvus_kb.py +++ b/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) diff --git a/docs/develop-guides/changelog.md b/docs/develop-guides/changelog.md index 3dc4524b..42f29119 100644 --- a/docs/develop-guides/changelog.md +++ b/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` 控制重排序和图检索参数展示。