35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
|
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"}
|