-
Notifications
You must be signed in to change notification settings - Fork 1.2k
feat: Implement keyword search in Qdrant #3099
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -4,7 +4,9 @@ | |
# This source code is licensed under the terms described in the LICENSE file in | ||
# the root directory of this source tree. | ||
|
||
import os | ||
import random | ||
import tempfile | ||
from unittest.mock import AsyncMock, MagicMock, patch | ||
|
||
import numpy as np | ||
|
@@ -18,7 +20,7 @@ | |
from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig | ||
from llama_stack.providers.inline.vector_io.faiss.faiss import FaissIndex, FaissVectorIOAdapter | ||
from llama_stack.providers.inline.vector_io.milvus.config import MilvusVectorIOConfig, SqliteKVStoreConfig | ||
from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig | ||
from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig | ||
from llama_stack.providers.inline.vector_io.sqlite_vec import SQLiteVectorIOConfig | ||
from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import SQLiteVecIndex, SQLiteVecVectorIOAdapter | ||
from llama_stack.providers.remote.vector_io.chroma.chroma import ChromaIndex, ChromaVectorIOAdapter, maybe_await | ||
|
@@ -32,7 +34,7 @@ | |
MILVUS_ALIAS = "test_milvus" | ||
|
||
|
||
@pytest.fixture(params=["milvus", "sqlite_vec", "faiss", "chroma", "pgvector"]) | ||
@pytest.fixture(params=["milvus", "sqlite_vec", "faiss", "chroma", "qdrant", "pgvector"]) | ||
def vector_provider(request): | ||
return request.param | ||
|
||
|
@@ -286,40 +288,53 @@ async def chroma_vec_adapter(chroma_vec_db_path, mock_inference_api, embedding_d | |
|
||
|
||
@pytest.fixture | ||
def qdrant_vec_db_path(tmp_path_factory): | ||
def qdrant_vec_db_path(tmp_path): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. why change from |
||
"""Use tmp_path with additional isolation to ensure unique path per test.""" | ||
import uuid | ||
|
||
db_path = str(tmp_path_factory.getbasetemp() / f"test_qdrant_{uuid.uuid4()}.db") | ||
return db_path | ||
# Create a completely isolated temporary directory | ||
temp_dir = tempfile.mkdtemp(prefix=f"qdrant_test_{uuid.uuid4()}_") | ||
return temp_dir | ||
|
||
|
||
@pytest.fixture | ||
async def qdrant_vec_adapter(qdrant_vec_db_path, mock_inference_api, embedding_dimension): | ||
import shutil | ||
import uuid | ||
|
||
config = QdrantVectorIOConfig( | ||
db_path=qdrant_vec_db_path, | ||
config = InlineQdrantVectorIOConfig( | ||
path=qdrant_vec_db_path, | ||
kvstore=SqliteKVStoreConfig(), | ||
) | ||
adapter = QdrantVectorIOAdapter( | ||
config=config, | ||
inference_api=mock_inference_api, | ||
files_api=None, | ||
) | ||
collection_id = f"qdrant_test_collection_{uuid.uuid4()}" | ||
|
||
original_initialize = adapter.initialize | ||
|
||
async def safe_initialize(): | ||
if not hasattr(adapter, "_initialized") or not adapter._initialized: | ||
await original_initialize() | ||
adapter._initialized = True | ||
|
||
adapter.initialize = safe_initialize | ||
await adapter.initialize() | ||
await adapter.register_vector_db( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. is this no longer needed? definitely right thing to do but just want to confirm. |
||
VectorDB( | ||
identifier=collection_id, | ||
provider_id="test_provider", | ||
embedding_model="test_model", | ||
embedding_dimension=embedding_dimension, | ||
) | ||
) | ||
|
||
collection_id = f"qdrant_test_collection_{uuid.uuid4()}" | ||
adapter.test_collection_id = collection_id | ||
adapter._test_db_path = qdrant_vec_db_path | ||
yield adapter | ||
|
||
await adapter.shutdown() | ||
|
||
try: | ||
if os.path.exists(qdrant_vec_db_path): | ||
shutil.rmtree(qdrant_vec_db_path, ignore_errors=True) | ||
except Exception: | ||
pass | ||
|
||
|
||
@pytest.fixture | ||
async def qdrant_vec_index(qdrant_vec_db_path, embedding_dimension): | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -145,3 +145,89 @@ async def test_qdrant_register_and_unregister_vector_db( | |
await qdrant_adapter.unregister_vector_db(vector_db_id) | ||
assert not (await qdrant_adapter.client.collection_exists(vector_db_id)) | ||
assert len((await qdrant_adapter.client.get_collections()).collections) == 0 | ||
|
||
|
||
# Keyword search tests | ||
async def test_query_chunks_keyword_search(qdrant_vec_index, sample_chunks, sample_embeddings): | ||
"""Test keyword search functionality in Qdrant.""" | ||
await qdrant_vec_index.add_chunks(sample_chunks, sample_embeddings) | ||
query_string = "Sentence 5" | ||
response = await qdrant_vec_index.query_keyword(query_string=query_string, k=3, score_threshold=0.0) | ||
|
||
assert isinstance(response, QueryChunksResponse) | ||
assert len(response.chunks) > 0, f"Expected some chunks, but got {len(response.chunks)}" | ||
|
||
non_existent_query_str = "blablabla" | ||
response_no_results = await qdrant_vec_index.query_keyword( | ||
query_string=non_existent_query_str, k=1, score_threshold=0.0 | ||
) | ||
|
||
assert isinstance(response_no_results, QueryChunksResponse) | ||
assert len(response_no_results.chunks) == 0, f"Expected 0 results, but got {len(response_no_results.chunks)}" | ||
|
||
|
||
async def test_query_chunks_keyword_search_k_greater_than_results(qdrant_vec_index, sample_chunks, sample_embeddings): | ||
"""Test keyword search when k is greater than available results.""" | ||
await qdrant_vec_index.add_chunks(sample_chunks, sample_embeddings) | ||
|
||
query_str = "Sentence 1 from document 0" # Should match only one chunk | ||
response = await qdrant_vec_index.query_keyword(k=5, score_threshold=0.0, query_string=query_str) | ||
|
||
assert isinstance(response, QueryChunksResponse) | ||
assert 0 < len(response.chunks) < 5, f"Expected results between [1, 4], got {len(response.chunks)}" | ||
assert any("Sentence 1 from document 0" in chunk.content for chunk in response.chunks), "Expected chunk not found" | ||
|
||
|
||
async def test_query_chunks_keyword_search_score_threshold(qdrant_vec_index, sample_chunks, sample_embeddings): | ||
"""Test keyword search with score threshold filtering.""" | ||
await qdrant_vec_index.add_chunks(sample_chunks, sample_embeddings) | ||
|
||
query_string = "Sentence 5" | ||
|
||
# Test with low threshold (should return results) | ||
response_low_threshold = await qdrant_vec_index.query_keyword(query_string=query_string, k=3, score_threshold=0.0) | ||
assert len(response_low_threshold.chunks) > 0 | ||
|
||
# Test with negative threshold (should return results since scores are 0.0) | ||
response_negative_threshold = await qdrant_vec_index.query_keyword( | ||
query_string=query_string, k=3, score_threshold=-1.0 | ||
) | ||
assert len(response_negative_threshold.chunks) > 0 | ||
|
||
|
||
async def test_query_chunks_keyword_search_edge_cases(qdrant_vec_index, sample_chunks, sample_embeddings): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. is this test actually needed? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't see why we shouldn't have it. This is lightweight and only testing that a valid There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'm intentionally not testing the return specifics, as that depends on the provider. But making sure that these values are at least accepted so that the input tests are at least robust. |
||
"""Test keyword search edge cases.""" | ||
await qdrant_vec_index.add_chunks(sample_chunks, sample_embeddings) | ||
|
||
# Test with empty string | ||
response_empty = await qdrant_vec_index.query_keyword(query_string="", k=3, score_threshold=0.0) | ||
assert isinstance(response_empty, QueryChunksResponse) | ||
|
||
# Test with very long query string | ||
long_query = "a" * 1000 | ||
response_long = await qdrant_vec_index.query_keyword(query_string=long_query, k=3, score_threshold=0.0) | ||
assert isinstance(response_long, QueryChunksResponse) | ||
|
||
# Test with special characters | ||
special_query = "!@#$%^&*()_+-=[]{}|;':\",./<>?" | ||
response_special = await qdrant_vec_index.query_keyword(query_string=special_query, k=3, score_threshold=0.0) | ||
assert isinstance(response_special, QueryChunksResponse) | ||
|
||
|
||
async def test_query_chunks_keyword_search_metadata_preservation( | ||
qdrant_vec_index, sample_chunks_with_metadata, sample_embeddings_with_metadata | ||
): | ||
"""Test that keyword search preserves chunk metadata.""" | ||
await qdrant_vec_index.add_chunks(sample_chunks_with_metadata, sample_embeddings_with_metadata) | ||
|
||
query_string = "Sentence 0" | ||
response = await qdrant_vec_index.query_keyword(query_string=query_string, k=2, score_threshold=0.0) | ||
|
||
assert len(response.chunks) > 0 | ||
for chunk in response.chunks: | ||
# Check that metadata is preserved | ||
assert hasattr(chunk, "metadata") or hasattr(chunk, "chunk_metadata") | ||
if hasattr(chunk, "chunk_metadata") and chunk.chunk_metadata: | ||
assert chunk.chunk_metadata.document_id is not None | ||
assert chunk.chunk_metadata.chunk_id is not None | ||
assert chunk.chunk_metadata.source is not None |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nice, looks like this is the appropriate way to handle things https://qdrant.tech/documentation/concepts/filtering/#full-text-match