326 lines
12 KiB
Python
326 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from yuxi.agents.backends.knowledge_base_backend import (
|
|
KBS_PATH,
|
|
KnowledgeBaseReadonlyBackend,
|
|
build_knowledge_base_filepath_map,
|
|
inject_filepaths_into_retrieval_result,
|
|
resolve_file_relative_virtual_path,
|
|
resolve_visible_knowledge_bases_for_context,
|
|
)
|
|
|
|
|
|
def test_knowledge_base_backend_builds_virtual_tree_and_materializes_files(monkeypatch, tmp_path) -> None:
|
|
visible_kbs = [
|
|
{"db_id": "db-1", "name": "FAQ"},
|
|
{"db_id": "db-2", "name": "FAQ"},
|
|
]
|
|
files_meta = {
|
|
"folder-api": {
|
|
"file_id": "folder-api",
|
|
"database_id": "db-1",
|
|
"parent_id": None,
|
|
"filename": "API",
|
|
"is_folder": True,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
"file-pdf": {
|
|
"file_id": "file-pdf",
|
|
"database_id": "db-1",
|
|
"parent_id": "folder-api",
|
|
"filename": "auth-guide.pdf",
|
|
"path": "http://minio/kb-source/db-1/upload/auth-guide.pdf",
|
|
"markdown_file": "http://minio/kb-parsed/db-1/parsed/file-pdf.md",
|
|
"size": 12,
|
|
"is_folder": False,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
"file-pdf-2": {
|
|
"file_id": "file-pdf-2",
|
|
"database_id": "db-1",
|
|
"parent_id": "folder-api",
|
|
"filename": "auth-guide.pdf",
|
|
"path": "http://minio/kb-source/db-1/upload/auth-guide-2.pdf",
|
|
"size": 7,
|
|
"is_folder": False,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
"file-db2": {
|
|
"file_id": "file-db2",
|
|
"database_id": "db-2",
|
|
"parent_id": None,
|
|
"filename": "overview.txt",
|
|
"path": "http://minio/kb-source/db-2/upload/overview.txt",
|
|
"size": 5,
|
|
"is_folder": False,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
}
|
|
monkeypatch.setattr("yuxi.agents.backends.knowledge_base_backend._all_files_meta", lambda: files_meta)
|
|
|
|
class _FakeMinio:
|
|
def download_file(self, bucket_name: str, object_name: str) -> bytes:
|
|
return f"{bucket_name}:{object_name}".encode()
|
|
|
|
monkeypatch.setattr("yuxi.agents.backends.knowledge_base_backend.get_minio_client", lambda: _FakeMinio())
|
|
|
|
backend = KnowledgeBaseReadonlyBackend(visible_kbs=visible_kbs, cache_root=tmp_path)
|
|
|
|
root_entries = {entry["path"] for entry in backend.ls_info("/")}
|
|
assert "/FAQ/" in root_entries
|
|
assert "/FAQ__db-2/" in root_entries
|
|
|
|
faq_entries = {entry["path"] for entry in backend.ls_info("/FAQ")}
|
|
assert "/FAQ/API/" in faq_entries
|
|
assert "/FAQ/parsed/" in faq_entries
|
|
|
|
api_entries = {entry["path"] for entry in backend.ls_info("/FAQ/API")}
|
|
assert "/FAQ/API/auth-guide.pdf" in api_entries
|
|
assert any(path.startswith("/FAQ/API/auth-guide.pdf__file-pdf") for path in api_entries)
|
|
|
|
parsed_entries = {entry["path"] for entry in backend.ls_info("/FAQ/parsed/API")}
|
|
assert "/FAQ/parsed/API/auth-guide.pdf.md" in parsed_entries
|
|
|
|
responses = backend.download_files(["/FAQ/API/auth-guide.pdf", "/FAQ/parsed/API/auth-guide.pdf.md"])
|
|
assert responses[0].content == b"kb-source:db-1/upload/auth-guide.pdf"
|
|
assert responses[1].content == b"kb-parsed:db-1/parsed/file-pdf.md"
|
|
|
|
assert "kb-source" in backend.read("/FAQ/API/auth-guide.pdf")
|
|
assert backend.write("/FAQ/API/new.txt", "x").error == "Knowledge base path is read-only."
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_visible_knowledge_bases_for_context_filters_by_enabled_names(monkeypatch) -> None:
|
|
async def _fake_get_databases_by_raw_id(user_id: int) -> dict:
|
|
assert user_id == 7
|
|
return {"databases": [{"db_id": "db-1", "name": "Alpha"}, {"db_id": "db-2", "name": "Beta"}]}
|
|
|
|
monkeypatch.setattr(
|
|
"yuxi.agents.backends.knowledge_base_backend.knowledge_base.get_databases_by_raw_id",
|
|
_fake_get_databases_by_raw_id,
|
|
)
|
|
|
|
context = SimpleNamespace(user_id="7", knowledges=["Beta"])
|
|
visible = await resolve_visible_knowledge_bases_for_context(context)
|
|
|
|
assert visible == [{"db_id": "db-2", "name": "Beta"}]
|
|
assert getattr(context, "_visible_knowledge_bases") == visible
|
|
|
|
|
|
def test_build_knowledge_base_filepath_map_matches_virtual_tree(monkeypatch, tmp_path) -> None:
|
|
visible_kbs = [
|
|
{"db_id": "db-1", "name": "FAQ"},
|
|
{"db_id": "db-2", "name": "FAQ"},
|
|
]
|
|
files_meta = {
|
|
"folder-api": {
|
|
"file_id": "folder-api",
|
|
"database_id": "db-1",
|
|
"parent_id": None,
|
|
"filename": "API",
|
|
"is_folder": True,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
"file-pdf": {
|
|
"file_id": "file-pdf",
|
|
"database_id": "db-1",
|
|
"parent_id": "folder-api",
|
|
"filename": "auth-guide.pdf",
|
|
"path": "http://minio/kb-source/db-1/upload/auth-guide.pdf",
|
|
"size": 12,
|
|
"is_folder": False,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
"file-pdf-2": {
|
|
"file_id": "file-pdf-2",
|
|
"database_id": "db-1",
|
|
"parent_id": "folder-api",
|
|
"filename": "auth-guide.pdf",
|
|
"path": "http://minio/kb-source/db-1/upload/auth-guide-2.pdf",
|
|
"size": 7,
|
|
"is_folder": False,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
"file-db2": {
|
|
"file_id": "file-db2",
|
|
"database_id": "db-2",
|
|
"parent_id": None,
|
|
"filename": "overview.txt",
|
|
"path": "http://minio/kb-source/db-2/upload/overview.txt",
|
|
"size": 5,
|
|
"is_folder": False,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
}
|
|
monkeypatch.setattr("yuxi.agents.backends.knowledge_base_backend._all_files_meta", lambda: files_meta)
|
|
|
|
filepath_map = build_knowledge_base_filepath_map(visible_kbs=visible_kbs, files_meta=files_meta)
|
|
backend = KnowledgeBaseReadonlyBackend(visible_kbs=visible_kbs, cache_root=tmp_path)
|
|
|
|
assert filepath_map["file-pdf"] == f"{KBS_PATH}/FAQ/API/auth-guide.pdf"
|
|
assert filepath_map["file-db2"] == f"{KBS_PATH}/FAQ__db-2/overview.txt"
|
|
assert filepath_map["file-pdf-2"].startswith(f"{KBS_PATH}/FAQ/API/auth-guide.pdf__")
|
|
|
|
mapped_virtual_paths = {path[len(KBS_PATH) :] for path in filepath_map.values()}
|
|
for virtual_path in mapped_virtual_paths:
|
|
assert virtual_path in backend._files
|
|
|
|
|
|
def test_resolve_file_relative_virtual_path_by_file_id(monkeypatch) -> None:
|
|
visible_kbs = [{"db_id": "db-1", "name": "食品 相关文献"}]
|
|
files_meta = {
|
|
"folder-1": {
|
|
"file_id": "folder-1",
|
|
"database_id": "db-1",
|
|
"parent_id": None,
|
|
"filename": "课程",
|
|
"is_folder": True,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
"file_79c496": {
|
|
"file_id": "file_79c496",
|
|
"database_id": "db-1",
|
|
"parent_id": "folder-1",
|
|
"filename": "营养与食品课程中的思政建设研究.pdf",
|
|
"path": "http://minio/knowledgebases/db-1/uploads/file_79c496.pdf",
|
|
"is_folder": False,
|
|
"size": 100,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
},
|
|
}
|
|
monkeypatch.setattr("yuxi.agents.backends.knowledge_base_backend._all_files_meta", lambda: files_meta)
|
|
|
|
relative_path = resolve_file_relative_virtual_path(file_id="file_79c496", visible_kbs=visible_kbs)
|
|
absolute_map = build_knowledge_base_filepath_map(visible_kbs=visible_kbs, files_meta=files_meta)
|
|
|
|
assert relative_path == "/食品 相关文献/课程/营养与食品课程中的思政建设研究.pdf"
|
|
assert absolute_map["file_79c496"] == f"{KBS_PATH}{relative_path}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_inject_filepaths_into_retrieval_result_injects_by_file_id(monkeypatch) -> None:
|
|
retrieval_chunks = [
|
|
{
|
|
"content": "auth guide",
|
|
"metadata": {
|
|
"file_id": "file-1",
|
|
"source": "auth-guide.pdf",
|
|
},
|
|
}
|
|
]
|
|
|
|
monkeypatch.setattr(
|
|
"yuxi.agents.backends.knowledge_base_backend._resolve_virtual_layout",
|
|
lambda **kwargs: SimpleNamespace(
|
|
source_filepaths={"file-1": f"{KBS_PATH}/FAQ/auth-guide.pdf"},
|
|
parsed_filepaths={"file-1": f"{KBS_PATH}/FAQ/parsed/auth-guide.pdf.md"},
|
|
),
|
|
)
|
|
|
|
result = await inject_filepaths_into_retrieval_result(
|
|
retrieval_chunks=retrieval_chunks,
|
|
visible_kbs=[{"db_id": "db-1", "name": "FAQ"}],
|
|
target_db_id="db-1",
|
|
)
|
|
|
|
assert result[0]["metadata"]["filepath"] == f"{KBS_PATH}/FAQ/auth-guide.pdf"
|
|
assert result[0]["metadata"]["parsed_path"] == f"{KBS_PATH}/FAQ/parsed/auth-guide.pdf.md"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_inject_filepaths_into_retrieval_result_injects_parsed_path_when_markdown_exists(monkeypatch) -> None:
|
|
files_meta = {
|
|
"file-1": {
|
|
"file_id": "file-1",
|
|
"database_id": "db-1",
|
|
"parent_id": None,
|
|
"filename": "auth-guide.pdf",
|
|
"path": "http://minio/kb-source/db-1/upload/auth-guide.pdf",
|
|
"markdown_file": "http://minio/kb-parsed/db-1/parsed/file-1.md",
|
|
"is_folder": False,
|
|
"created_at": "2026-03-26T00:00:00Z",
|
|
}
|
|
}
|
|
monkeypatch.setattr("yuxi.agents.backends.knowledge_base_backend._all_files_meta", lambda: files_meta)
|
|
|
|
retrieval_chunks = [{"content": "auth guide", "metadata": {"file_id": "file-1"}}]
|
|
result = await inject_filepaths_into_retrieval_result(
|
|
retrieval_chunks=retrieval_chunks,
|
|
visible_kbs=[{"db_id": "db-1", "name": "FAQ"}],
|
|
target_db_id="db-1",
|
|
)
|
|
|
|
assert result[0]["metadata"]["filepath"] == f"{KBS_PATH}/FAQ/auth-guide.pdf"
|
|
assert result[0]["metadata"]["parsed_path"] == f"{KBS_PATH}/FAQ/parsed/auth-guide.pdf.md"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_inject_filepaths_into_retrieval_result_does_not_use_filename_fallback(monkeypatch) -> None:
|
|
retrieval_chunks = [
|
|
{
|
|
"id": "chunk-1",
|
|
"metadata": {
|
|
"source": "auth-guide.pdf",
|
|
},
|
|
}
|
|
]
|
|
|
|
monkeypatch.setattr(
|
|
"yuxi.agents.backends.knowledge_base_backend._resolve_virtual_layout",
|
|
lambda **kwargs: SimpleNamespace(
|
|
source_filepaths={
|
|
"file-1": f"{KBS_PATH}/FAQ/API/auth-guide.pdf",
|
|
"file-2": f"{KBS_PATH}/FAQ/API/another.pdf",
|
|
},
|
|
parsed_filepaths={
|
|
"file-1": f"{KBS_PATH}/FAQ/parsed/API/auth-guide.pdf.md",
|
|
"file-2": f"{KBS_PATH}/FAQ/parsed/API/another.pdf.md",
|
|
},
|
|
),
|
|
)
|
|
|
|
result = await inject_filepaths_into_retrieval_result(
|
|
retrieval_chunks=retrieval_chunks,
|
|
visible_kbs=[{"db_id": "db-1", "name": "FAQ"}],
|
|
target_db_id="db-1",
|
|
)
|
|
|
|
assert "filepath" not in result[0]["metadata"]
|
|
assert "parsed_path" not in result[0]["metadata"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_inject_filepaths_into_retrieval_result_requires_explicit_file_id(monkeypatch) -> None:
|
|
retrieval_chunks = [
|
|
{
|
|
"content": "chunk",
|
|
"metadata": {
|
|
"source": "http://172.19.13.5:9000/knowledgebases/kb-1/parsed/file_79c496.md",
|
|
},
|
|
}
|
|
]
|
|
|
|
monkeypatch.setattr(
|
|
"yuxi.agents.backends.knowledge_base_backend._resolve_virtual_layout",
|
|
lambda **kwargs: SimpleNamespace(
|
|
source_filepaths={
|
|
"file_79c496": f"{KBS_PATH}/食品 相关文献/课程/营养与食品课程中的思政建设研究.pdf",
|
|
},
|
|
parsed_filepaths={
|
|
"file_79c496": f"{KBS_PATH}/食品 相关文献/parsed/课程/营养与食品课程中的思政建设研究.pdf.md",
|
|
},
|
|
),
|
|
)
|
|
|
|
result = await inject_filepaths_into_retrieval_result(
|
|
retrieval_chunks=retrieval_chunks,
|
|
visible_kbs=[{"db_id": "db-1", "name": "食品 相关文献"}],
|
|
target_db_id="db-1",
|
|
)
|
|
|
|
assert "filepath" not in result[0]["metadata"]
|
|
assert "parsed_path" not in result[0]["metadata"]
|