feat(kb): 添加打开知识库文档功能及相关输入模型，优化文件内容读取
This commit is contained in:
Wenjie Zhang
parent
fd9b770ce9
commit
3771b7e804
@ -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"]
|
||||
|
||||
@ -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]
|
||||
|
||||
@ -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:
|
||||
"""
|
||||
|
||||
@ -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()
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
34
backend/test/unit/knowledge/test_lightrag_chunk_mapping.py
Normal file
34
backend/test/unit/knowledge/test_lightrag_chunk_mapping.py
Normal file
@ -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"}
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user