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

feat (RAG): Implement configurable search mode in RAGQueryConfig

Browse source 
Signed-off-by: Varsha Prasad Narsing <varshaprasad96@gmail.com>
This commit is contained in:
Varsha Prasad Narsing 2025-04-14 16:53:17 -07:00
parent 85b5f3172b
commit e2a7022d3c
14 changed files with 210 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

40
tests/unit/providers/vector_io/test_sqlite_vec.py
View file

@ -57,14 +57,50 @@ async def test_add_chunks(sqlite_vec_index, sample_chunks, sample_embeddings):
@pytest.mark.asyncio
async def test_query_chunks(sqlite_vec_index, sample_chunks, sample_embeddings, embedding_dimension):
async def test_query_chunks_vector(sqlite_vec_index, sample_chunks, sample_embeddings, embedding_dimension):
await sqlite_vec_index.add_chunks(sample_chunks, sample_embeddings)
query_embedding = np.random.rand(embedding_dimension).astype(np.float32)
response = await sqlite_vec_index.query(query_embedding, k=2, score_threshold=0.0)
response = await sqlite_vec_index.query(query_embedding, query_string="", k=2, score_threshold=0.0, mode="vector")
assert isinstance(response, QueryChunksResponse)
assert len(response.chunks) == 2
@pytest.mark.asyncio
async def test_query_chunks_full_text_search(sqlite_vec_index, sample_chunks, sample_embeddings):
await sqlite_vec_index.add_chunks(sample_chunks, sample_embeddings)
query_string = "Sentence 5"
response = await sqlite_vec_index.query(
embedding=None, k=3, score_threshold=0.0, query_string=query_string, mode="keyword"
)
assert isinstance(response, QueryChunksResponse)
assert len(response.chunks) == 3, f"Expected at least one result, but got {len(response.chunks)}"
non_existent_query_str = "blablabla"
response_no_results = await sqlite_vec_index.query(
embedding=None, query_string=non_existent_query_str, k=1, score_threshold=0.0, mode="keyword"
)
assert isinstance(response_no_results, QueryChunksResponse)
assert len(response_no_results.chunks) == 0, f"Expected 0 results, but got {len(response_no_results.chunks)}"
@pytest.mark.asyncio
async def test_query_chunks_full_text_search_k_greater_than_results(sqlite_vec_index, sample_chunks, sample_embeddings):
# Re-initialize with a clean index
await sqlite_vec_index.add_chunks(sample_chunks, sample_embeddings)
query_str = "Sentence 1 from document 0" # Should match only one chunk
response = await sqlite_vec_index.query(
embedding=None, k=5, score_threshold=0.0, query_string=query_str, mode="keyword"
)
assert isinstance(response, QueryChunksResponse)
assert 0 < len(response.chunks) < 5, f"Expected <5 results but >0, got {len(response.chunks)}"
assert any("Sentence 1 from document 0" in chunk.content for chunk in response.chunks), "Expected chunk not found"
@pytest.mark.asyncio
async def test_chunk_id_conflict(sqlite_vec_index, sample_chunks, embedding_dimension):
"""Test that chunk IDs do not conflict across batches when inserting chunks."""