phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-05 10:23:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge 9aef325934 into sapling-pr-archive-ehhuang

Browse source
This commit is contained in:
ehhuang 2025-10-27 15:32:50 -07:00 • committed by GitHub
commit e9a8967ed5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
41 changed files with 1280 additions and 197 deletions
Download patch file Download diff file Expand all files Collapse all files

18
tests/unit/rag/test_rag_query.py
View file

@ -23,14 +23,14 @@ class TestRagQuery:
config=MagicMock(), vector_io_api=MagicMock(), inference_api=MagicMock(), files_api=MagicMock()
)
with pytest.raises(ValueError):
await rag_tool.query(content=MagicMock(), vector_db_ids=[])
await rag_tool.query(content=MagicMock(), vector_store_ids=[])
async def test_query_chunk_metadata_handling(self):
rag_tool = MemoryToolRuntimeImpl(
config=MagicMock(), vector_io_api=MagicMock(), inference_api=MagicMock(), files_api=MagicMock()
)
content = "test query content"
vector_db_ids = ["db1"]
vector_store_ids = ["db1"]
chunk_metadata = ChunkMetadata(
document_id="doc1",
@ -55,7 +55,7 @@ class TestRagQuery:
query_response = QueryChunksResponse(chunks=[chunk], scores=[1.0])
rag_tool.vector_io_api.query_chunks = AsyncMock(return_value=query_response)
result = await rag_tool.query(content=content, vector_db_ids=vector_db_ids)
result = await rag_tool.query(content=content, vector_store_ids=vector_store_ids)
assert result is not None
expected_metadata_string = (
@ -90,7 +90,7 @@ class TestRagQuery:
files_api=MagicMock(),
)
vector_db_ids = ["db1", "db2"]
vector_store_ids = ["db1", "db2"]
# Fake chunks from each DB
chunk_metadata1 = ChunkMetadata(
@ -101,7 +101,7 @@ class TestRagQuery:
)
chunk1 = Chunk(
content="chunk from db1",
metadata={"vector_db_id": "db1", "document_id": "doc1"},
metadata={"vector_store_id": "db1", "document_id": "doc1"},
stored_chunk_id="c1",
chunk_metadata=chunk_metadata1,
)
@ -114,7 +114,7 @@ class TestRagQuery:
)
chunk2 = Chunk(
content="chunk from db2",
metadata={"vector_db_id": "db2", "document_id": "doc2"},
metadata={"vector_store_id": "db2", "document_id": "doc2"},
stored_chunk_id="c2",
chunk_metadata=chunk_metadata2,
)
@ -126,13 +126,13 @@ class TestRagQuery:
]
)
result = await rag_tool.query(content="test", vector_db_ids=vector_db_ids)
result = await rag_tool.query(content="test", vector_store_ids=vector_store_ids)
returned_chunks = result.metadata["chunks"]
returned_scores = result.metadata["scores"]
returned_doc_ids = result.metadata["document_ids"]
returned_vector_db_ids = result.metadata["vector_db_ids"]
returned_vector_store_ids = result.metadata["vector_store_ids"]
assert returned_chunks == ["chunk from db1", "chunk from db2"]
assert returned_scores == (0.9, 0.8)
assert returned_doc_ids == ["doc1", "doc2"]
assert returned_vector_db_ids == ["db1", "db2"]
assert returned_vector_store_ids == ["db1", "db2"]