diff --git a/backend/package/yuxi/agents/toolkits/kbs/__init__.py b/backend/package/yuxi/agents/toolkits/kbs/__init__.py index d771ed2b..dc052ee1 100644 --- a/backend/package/yuxi/agents/toolkits/kbs/__init__.py +++ b/backend/package/yuxi/agents/toolkits/kbs/__init__.py @@ -1,5 +1,6 @@ from .tools import ( get_common_kb_tools, + open_kb_document, ) -__all__ = ["get_common_kb_tools"] +__all__ = ["get_common_kb_tools", "open_kb_document"] diff --git a/backend/package/yuxi/agents/toolkits/kbs/tools.py b/backend/package/yuxi/agents/toolkits/kbs/tools.py index e572ba4e..02ee8a46 100644 --- a/backend/package/yuxi/agents/toolkits/kbs/tools.py +++ b/backend/package/yuxi/agents/toolkits/kbs/tools.py @@ -161,6 +161,16 @@ class QueryKBInput(BaseModel): ) +class OpenKBDocumentInput(BaseModel): + """打开知识库文档输入模型""" + + resource_id: str = Field(description="知识库资源 ID,当前对应知识库 db_id") + file_id: str = Field(description="要打开的文档 ID,也就是知识库文件 file_id") + line: int | None = Field(default=None, ge=1, description="可选,1-based 起始行号") + offset: int | None = Field(default=None, ge=0, description="可选,0-based 起始偏移;line 优先于 offset") + window_size: int = Field(default=800, ge=1, le=2000, description="读取窗口行数,默认 800 行") + + async def _resolve_visible_knowledge_bases_for_query(runtime: ToolRuntime | None) -> list[dict[str, Any]]: if runtime is None: return [] @@ -188,31 +198,54 @@ def _find_query_target( retrievers: dict[str, Any], visible_kbs: list[dict[str, Any]], ) -> tuple[str | None, dict[str, Any] | None, str | None]: - if visible_kbs: - matched_kbs = [db for db in visible_kbs if str(db.get("name") or "").strip() == kb_name] - if not matched_kbs: - return None, None, f"知识库 '{kb_name}' 不存在或当前会话未启用" - if len(matched_kbs) > 1: - return None, None, f"知识库 '{kb_name}' 存在重名,请先调整名称后重试" + if not visible_kbs: + return None, None, "无法获取当前会话可访问的知识库" - target_db_id = str(matched_kbs[0].get("db_id") or "") - target_info = retrievers.get(target_db_id) - if target_info is None: - return None, None, f"知识库 '{kb_name}' 不存在" - return target_db_id, target_info, None + matched_kbs = [db for db in visible_kbs if str(db.get("name") or "").strip() == kb_name] + if not matched_kbs: + return None, None, f"知识库 '{kb_name}' 不存在或当前会话未启用" + if len(matched_kbs) > 1: + return None, None, f"知识库 '{kb_name}' 存在重名,请先调整名称后重试" - for db_id, info in retrievers.items(): - if info["name"] == kb_name: - return str(db_id), info, None + target_db_id = str(matched_kbs[0].get("db_id") or "") + target_info = retrievers.get(target_db_id) + if target_info is None: + return None, None, f"知识库 '{kb_name}' 不存在" + return target_db_id, target_info, None - return None, None, f"知识库 '{kb_name}' 不存在" + +def _normalize_retrieval_result_metadata(result: Any) -> Any: + if not isinstance(result, list): + return result + + for chunk in result: + if not isinstance(chunk, dict): + continue + + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + + file_id = metadata.get("file_id") or chunk.get("file_id") or chunk.get("full_doc_id") + if file_id and not metadata.get("file_id"): + metadata["file_id"] = str(file_id) + + for key in ("chunk_id", "chunk_index"): + value = metadata.get(key) if metadata.get(key) is not None else chunk.get(key) + if value is not None and metadata.get(key) is None: + metadata[key] = value + + chunk["metadata"] = metadata + + return result @tool(args_schema=QueryKBInput) async def query_kb(kb_name: str, query_text: str, file_name: str | None = None, runtime: ToolRuntime = None) -> Any: """在指定知识库中检索内容 - 当用户需要查询具体内容时使用此工具。根据关键词在知识库中检索相关文档片段。 + 当用户需要查询具体内容时使用此工具。根据关键词在知识库中检索相关文档片段;结果中的 + metadata.file_id 和知识库 resource_id 可用于继续调用 open_kb_document 打开原文。 Args: kb_name: 知识库名称 @@ -254,6 +287,8 @@ async def query_kb(kb_name: str, query_text: str, file_name: str | None = None, else: result = retriever(query_text, **kwargs) + result = _normalize_retrieval_result_metadata(result) + if kb_type != "milvus": return result @@ -275,12 +310,67 @@ async def query_kb(kb_name: str, query_text: str, file_name: str | None = None, return f"检索失败: {str(e)}" +@tool(args_schema=OpenKBDocumentInput) +async def open_kb_document( + resource_id: str, + file_id: str, + line: int | None = None, + offset: int | None = None, + window_size: int = 800, + runtime: ToolRuntime = None, +) -> dict[str, Any] | str: + """按行窗口打开知识库文档原文 + + 当 query_kb 返回的片段不足以回答问题,或需要查看某个文档的上下文时使用。 + resource_id 对应知识库 db_id,file_id 对应检索结果 metadata.file_id。 + """ + normalized_resource_id = str(resource_id or "").strip() + normalized_file_id = str(file_id or "").strip() + if not normalized_resource_id: + return "请提供 resource_id" + if not normalized_file_id: + return "请提供 file_id" + + visible_kbs = await _resolve_visible_knowledge_bases_for_query(runtime) + if not visible_kbs: + return "无法获取当前会话可访问的知识库" + + visible_resource_ids = {str(kb.get("db_id") or "").strip() for kb in visible_kbs} + if normalized_resource_id not in visible_resource_ids: + return f"知识库资源 '{normalized_resource_id}' 不存在或当前会话未启用" + + retrievers = knowledge_base.get_retrievers() + target_info = retrievers.get(normalized_resource_id) + if target_info is None: + return f"知识库资源 '{normalized_resource_id}' 不存在" + + metadata = target_info.get("metadata") if isinstance(target_info, dict) else None + kb_type = str((metadata or {}).get("kb_type") or "").strip().lower() + if kb_type == "dify": + return "Dify 知识库为外部只读检索源,当前不支持通过 Open 打开全文" + + try: + start_offset = int(line) - 1 if line is not None else int(offset or 0) + window = await knowledge_base.open_file_content( + normalized_resource_id, + normalized_file_id, + offset=start_offset, + limit=window_size, + ) + return {"resource_id": normalized_resource_id, "file_id": normalized_file_id, **window} + + except Exception as e: + logger.error(f"打开知识库文档失败: {e}") + return f"打开知识库文档失败: {str(e)}" + + def get_common_kb_tools() -> list: """获取通用知识库工具列表 - 返回 3 个通用工具: + 返回 4 个通用工具: - list_kbs: 列出用户可访问的知识库 - get_mindmap: 获取指定知识库的思维导图 - query_kb: 在指定知识库中检索 + - open_kb_document: 按 file_id 分段打开知识库文档 """ - return [list_kbs, get_mindmap, query_kb] + return [list_kbs, get_mindmap, query_kb, open_kb_document] diff --git a/backend/package/yuxi/knowledge/base.py b/backend/package/yuxi/knowledge/base.py index 0ccbf048..193538cd 100644 --- a/backend/package/yuxi/knowledge/base.py +++ b/backend/package/yuxi/knowledge/base.py @@ -397,6 +397,43 @@ class KnowledgeBase(ABC): content_bytes = await minio_client.adownload_file(bucket_name, object_name) return content_bytes.decode("utf-8") + def _build_open_file_window(self, content: str, *, offset: int = 0, limit: int = 800) -> dict[str, Any]: + lines = content.splitlines() + total_lines = len(lines) + start = min(max(int(offset), 0), total_lines) + window_size = min(max(int(limit), 1), 2000) + selected = lines[start : start + window_size] + end = start + len(selected) + + return { + "start_line": start + 1 if selected else 0, + "end_line": end, + "total_lines": total_lines, + "offset": start, + "window_size": window_size, + "has_more_before": start > 0, + "has_more_after": end < total_lines, + "next_offset": end if end < total_lines else None, + "content": "\n".join(f"{start + idx + 1:6d}\t{line}" for idx, line in enumerate(selected)), + } + + async def open_file_content(self, db_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict: + """按行窗口打开文件解析后的 Markdown 内容""" + file_meta = self.files_meta.get(file_id) + if file_meta is None: + raise Exception(f"文件不存在: {file_id}") + if file_meta.get("database_id") != db_id: + raise Exception(f"文件 {file_id} 不属于知识库 {db_id}") + if file_meta.get("is_folder"): + raise Exception(f"文件 {file_id} 是文件夹") + + markdown_file = file_meta.get("markdown_file") + if not markdown_file: + raise Exception(f"文件 {file_id} 没有解析后的 Markdown 内容") + + content = await self._read_markdown_from_minio(markdown_file) + return self._build_open_file_window(content, offset=offset, limit=limit) + @abstractmethod async def index_file(self, db_id: str, file_id: str, operator_id: str | None = None) -> dict: """ diff --git a/backend/package/yuxi/knowledge/implementations/dify.py b/backend/package/yuxi/knowledge/implementations/dify.py index e95551c5..09fb1b22 100644 --- a/backend/package/yuxi/knowledge/implementations/dify.py +++ b/backend/package/yuxi/knowledge/implementations/dify.py @@ -63,6 +63,10 @@ class DifyKB(KnowledgeBase): async def get_file_content(self, db_id: str, file_id: str) -> dict: raise self._readonly_error() + async def open_file_content(self, db_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict: + del offset, limit + raise self._readonly_error() + async def get_file_info(self, db_id: str, file_id: str) -> dict: raise self._readonly_error() diff --git a/backend/package/yuxi/knowledge/implementations/lightrag.py b/backend/package/yuxi/knowledge/implementations/lightrag.py index 9e135c0f..815968cf 100644 --- a/backend/package/yuxi/knowledge/implementations/lightrag.py +++ b/backend/package/yuxi/knowledge/implementations/lightrag.py @@ -528,6 +528,36 @@ class LightRagKB(KnowledgeBase): return processed_items_info + @staticmethod + def _attach_file_ids_to_chunks(rag: LightRAG, chunks): + if not isinstance(chunks, list): + return chunks + + text_chunks = getattr(rag, "text_chunks", None) + chunk_store = getattr(text_chunks, "_data", None) + if not isinstance(chunk_store, dict): + return chunks + + for chunk in chunks: + if not isinstance(chunk, dict): + continue + + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + + chunk_id = metadata.get("chunk_id") or chunk.get("chunk_id") or chunk.get("id") + if chunk_id and not metadata.get("chunk_id"): + metadata["chunk_id"] = str(chunk_id) + + stored_chunk = chunk_store.get(chunk_id) + if isinstance(stored_chunk, dict) and stored_chunk.get("full_doc_id") and not metadata.get("file_id"): + metadata["file_id"] = str(stored_chunk["full_doc_id"]) + + chunk["metadata"] = metadata + + return chunks + async def aquery(self, query_text: str, db_id: str, agent_call: bool = False, **kwargs) -> str: """异步查询知识库""" rag = await self._get_lightrag_instance(db_id) @@ -579,7 +609,7 @@ class LightRagKB(KnowledgeBase): data = response.get("data", {}) or {} if scope == "chunks": - return data.get("chunks", []) + return self._attach_file_ids_to_chunks(rag, data.get("chunks", [])) result = {} if scope in ["graph", "all"]: @@ -594,7 +624,7 @@ class LightRagKB(KnowledgeBase): result["references"] = data.get("references", []) if scope == "all": - result["chunks"] = data.get("chunks", []) + result["chunks"] = self._attach_file_ids_to_chunks(rag, data.get("chunks", [])) return result diff --git a/backend/package/yuxi/knowledge/manager.py b/backend/package/yuxi/knowledge/manager.py index 23f8b3e6..3dc2b4ad 100644 --- a/backend/package/yuxi/knowledge/manager.py +++ b/backend/package/yuxi/knowledge/manager.py @@ -505,6 +505,11 @@ class KnowledgeBaseManager: kb_instance = await self._get_kb_for_database(db_id) return await kb_instance.get_file_content(db_id, file_id) + async def open_file_content(self, db_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict: + """按行窗口打开文件解析后的 Markdown 内容""" + kb_instance = await self._get_kb_for_database(db_id) + return await kb_instance.open_file_content(db_id, file_id, offset, limit) + async def get_file_info(self, db_id: str, file_id: str) -> dict: """获取文件完整信息(基本信息+内容信息)- 保持向后兼容""" kb_instance = await self._get_kb_for_database(db_id) diff --git a/backend/test/unit/knowledge/test_lightrag_chunk_mapping.py b/backend/test/unit/knowledge/test_lightrag_chunk_mapping.py new file mode 100644 index 00000000..d907249c --- /dev/null +++ b/backend/test/unit/knowledge/test_lightrag_chunk_mapping.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from yuxi.knowledge.implementations.lightrag import LightRagKB + + +def test_lightrag_attaches_file_id_to_chunks_from_chunk_store() -> None: + rag = SimpleNamespace( + text_chunks=SimpleNamespace( + _data={ + "chunk-1": {"full_doc_id": "file-1"}, + "chunk-2": {"full_doc_id": "file-2"}, + } + ) + ) + chunks = [ + {"chunk_id": "chunk-1", "content": "one"}, + {"id": "chunk-2", "content": "two", "metadata": {}}, + ] + + result = LightRagKB._attach_file_ids_to_chunks(rag, chunks) + + assert result[0]["metadata"] == {"chunk_id": "chunk-1", "file_id": "file-1"} + assert result[1]["metadata"] == {"chunk_id": "chunk-2", "file_id": "file-2"} + + +def test_lightrag_keeps_existing_file_id_when_mapping_chunks() -> None: + rag = SimpleNamespace(text_chunks=SimpleNamespace(_data={"chunk-1": {"full_doc_id": "file-1"}})) + chunks = [{"metadata": {"chunk_id": "chunk-1", "file_id": "existing-file"}}] + + result = LightRagKB._attach_file_ids_to_chunks(rag, chunks) + + assert result[0]["metadata"] == {"chunk_id": "chunk-1", "file_id": "existing-file"} diff --git a/backend/test/unit/toolkits/test_kbs_tools.py b/backend/test/unit/toolkits/test_kbs_tools.py index 888c8be0..b6851063 100644 --- a/backend/test/unit/toolkits/test_kbs_tools.py +++ b/backend/test/unit/toolkits/test_kbs_tools.py @@ -8,26 +8,59 @@ import pytest from yuxi.agents.toolkits.kbs import tools +def _tool_callable(tool): + callback = getattr(tool, "coroutine", None) + if callback is not None: + return callback + + callback = getattr(tool, "func", None) + if callback is not None: + return callback + + raise AssertionError(f"{tool.name} tool has no callable entry") + + def _query_kb_callable(): - callback = getattr(tools.query_kb, "coroutine", None) - if callback is not None: - return callback - - callback = getattr(tools.query_kb, "func", None) - if callback is not None: - return callback - - raise AssertionError("query_kb tool has no callable entry") + return _tool_callable(tools.query_kb) -async def _run_query_kb(**kwargs): - callback = _query_kb_callable() +def _open_kb_document_callable(): + return _tool_callable(tools.open_kb_document) + + +async def _run_tool(callback, **kwargs): result = callback(**kwargs) if inspect.isawaitable(result): return await result return result +async def _run_query_kb(**kwargs): + return await _run_tool(_query_kb_callable(), **kwargs) + + +async def _run_open_kb_document(**kwargs): + return await _run_tool(_open_kb_document_callable(), **kwargs) + + +def _build_test_window(content: str, offset: int = 0, limit: int = 800) -> dict: + lines = content.splitlines() + start = min(max(offset, 0), len(lines)) + selected = lines[start : start + limit] + end = start + len(selected) + return { + "start_line": start + 1 if selected else 0, + "end_line": end, + "total_lines": len(lines), + "offset": start, + "window_size": limit, + "has_more_before": start > 0, + "has_more_after": end < len(lines), + "next_offset": end if end < len(lines) else None, + "content": "\n".join(f"{start + idx + 1:6d}\t{line}" for idx, line in enumerate(selected)), + } + + @pytest.mark.asyncio async def test_query_kb_injects_filepath_into_chunk_metadata(monkeypatch) -> None: async def _fake_retriever(query_text: str, **kwargs): @@ -209,3 +242,175 @@ async def test_query_kb_uses_backend_filepath_injector(monkeypatch) -> None: assert result[0]["metadata"]["filepath"] == "/home/gem/kbs/FAQ/auth-guide.pdf" assert result[0]["metadata"]["parsed_path"] == "/home/gem/kbs/FAQ/parsed/auth-guide.pdf.md" + + +@pytest.mark.asyncio +async def test_query_kb_normalizes_file_metadata_for_open(monkeypatch) -> None: + async def _fake_retriever(query_text: str, **kwargs): + assert query_text == "auth" + return [ + { + "content": "auth guide", + "full_doc_id": "file-1", + "chunk_id": "chunk-1", + "chunk_index": 3, + } + ] + + monkeypatch.setattr( + tools.knowledge_base, + "get_retrievers", + lambda: { + "db-1": { + "name": "FAQ", + "retriever": _fake_retriever, + "metadata": {"kb_type": "lightrag"}, + } + }, + ) + + async def _fake_visible_kbs(runtime): + return [{"db_id": "db-1", "name": "FAQ"}] + + monkeypatch.setattr(tools, "_resolve_visible_knowledge_bases_for_query", _fake_visible_kbs) + + runtime = SimpleNamespace(context=SimpleNamespace()) + result = await _run_query_kb(kb_name="FAQ", query_text="auth", runtime=runtime) + + assert result[0]["metadata"] == { + "file_id": "file-1", + "chunk_id": "chunk-1", + "chunk_index": 3, + } + + +@pytest.mark.asyncio +async def test_open_kb_document_reads_markdown_content_by_default_window(monkeypatch) -> None: + lines = [f"line {index}" for index in range(1, 1001)] + + monkeypatch.setattr( + tools.knowledge_base, + "get_retrievers", + lambda: { + "db-1": { + "name": "FAQ", + "retriever": object(), + "metadata": {"kb_type": "milvus"}, + } + }, + ) + + async def _fake_visible_kbs(runtime): + return [{"db_id": "db-1", "name": "FAQ"}] + + async def _fake_open_file_content(db_id: str, file_id: str, offset: int = 0, limit: int = 800): + assert db_id == "db-1" + assert file_id == "file-1" + return _build_test_window("\n".join(lines), offset=offset, limit=limit) + + monkeypatch.setattr(tools, "_resolve_visible_knowledge_bases_for_query", _fake_visible_kbs) + monkeypatch.setattr(tools.knowledge_base, "open_file_content", _fake_open_file_content) + + runtime = SimpleNamespace(context=SimpleNamespace()) + result = await _run_open_kb_document(resource_id="db-1", file_id="file-1", runtime=runtime) + + assert result["resource_id"] == "db-1" + assert result["file_id"] == "file-1" + assert result["start_line"] == 1 + assert result["end_line"] == 800 + assert result["total_lines"] == 1000 + assert result["window_size"] == 800 + assert result["has_more_before"] is False + assert result["has_more_after"] is True + assert result["next_offset"] == 800 + assert " 1\tline 1" in result["content"] + assert " 800\tline 800" in result["content"] + + +@pytest.mark.asyncio +async def test_open_kb_document_prefers_line_over_offset(monkeypatch) -> None: + lines = [f"line {index}" for index in range(1, 1001)] + + monkeypatch.setattr( + tools.knowledge_base, + "get_retrievers", + lambda: { + "db-1": { + "name": "FAQ", + "retriever": object(), + "metadata": {"kb_type": "milvus"}, + } + }, + ) + + async def _fake_visible_kbs(runtime): + return [{"db_id": "db-1", "name": "FAQ"}] + + async def _fake_open_file_content(db_id: str, file_id: str, offset: int = 0, limit: int = 800): + assert db_id == "db-1" + assert file_id == "file-1" + return _build_test_window("\n".join(lines), offset=offset, limit=limit) + + monkeypatch.setattr(tools, "_resolve_visible_knowledge_bases_for_query", _fake_visible_kbs) + monkeypatch.setattr(tools.knowledge_base, "open_file_content", _fake_open_file_content) + + runtime = SimpleNamespace(context=SimpleNamespace()) + result = await _run_open_kb_document( + resource_id="db-1", + file_id="file-1", + line=801, + offset=0, + window_size=10, + runtime=runtime, + ) + + assert result["offset"] == 800 + assert result["start_line"] == 801 + assert result["end_line"] == 810 + assert result["has_more_before"] is True + assert result["has_more_after"] is True + assert result["next_offset"] == 810 + assert " 801\tline 801" in result["content"] + + +@pytest.mark.asyncio +async def test_open_kb_document_rejects_invisible_resource(monkeypatch) -> None: + async def _fake_visible_kbs(runtime): + return [{"db_id": "db-2", "name": "FAQ"}] + + monkeypatch.setattr(tools, "_resolve_visible_knowledge_bases_for_query", _fake_visible_kbs) + + runtime = SimpleNamespace(context=SimpleNamespace()) + result = await _run_open_kb_document(resource_id="db-1", file_id="file-1", runtime=runtime) + + assert "不存在或当前会话未启用" in result + + +@pytest.mark.asyncio +async def test_open_kb_document_requires_markdown_content(monkeypatch) -> None: + monkeypatch.setattr( + tools.knowledge_base, + "get_retrievers", + lambda: { + "db-1": { + "name": "FAQ", + "retriever": object(), + "metadata": {"kb_type": "milvus"}, + } + }, + ) + + async def _fake_visible_kbs(runtime): + return [{"db_id": "db-1", "name": "FAQ"}] + + async def _fake_open_file_content(db_id: str, file_id: str, offset: int = 0, limit: int = 800): + del db_id, file_id, offset, limit + raise Exception("文件 file-1 没有解析后的 Markdown 内容") + + monkeypatch.setattr(tools, "_resolve_visible_knowledge_bases_for_query", _fake_visible_kbs) + monkeypatch.setattr(tools.knowledge_base, "open_file_content", _fake_open_file_content) + + runtime = SimpleNamespace(context=SimpleNamespace()) + result = await _run_open_kb_document(resource_id="db-1", file_id="file-1", runtime=runtime) + + assert "没有解析后的 Markdown 内容" in result