Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 37 additions & 7 deletions llama_stack/providers/remote/vector_io/qdrant/qdrant.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,13 +128,43 @@ async def query_vector(self, embedding: NDArray, k: int, score_threshold: float)

return QueryChunksResponse(chunks=chunks, scores=scores)

async def query_keyword(
self,
query_string: str,
k: int,
score_threshold: float,
) -> QueryChunksResponse:
raise NotImplementedError("Keyword search is not supported in Qdrant")
async def query_keyword(self, query_string: str, k: int, score_threshold: float) -> QueryChunksResponse:
try:
results = (
await self.client.query_points(
collection_name=self.collection_name,
query_filter=models.Filter(
Copy link
Collaborator

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

must=[
models.FieldCondition(
key="chunk_content.content", match=models.MatchText(text=query_string)
)
]
),
limit=k,
with_payload=True,
with_vectors=False,
score_threshold=score_threshold,
)
).points
except Exception as e:
log.error(f"Error querying keyword search in Qdrant collection {self.collection_name}: {e}")
raise

chunks, scores = [], []
for point in results:
assert isinstance(point, models.ScoredPoint)
assert point.payload is not None

try:
chunk = Chunk(**point.payload["chunk_content"])
except Exception:
log.exception("Failed to parse chunk")
continue

chunks.append(chunk)
scores.append(point.score)

return QueryChunksResponse(chunks=chunks, scores=scores)

async def query_hybrid(
self,
Expand Down
2 changes: 2 additions & 0 deletions tests/integration/vector_io/test_openai_vector_stores.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ def skip_if_provider_doesnt_support_openai_vector_stores_search(client_with_mode
],
"keyword": [
"inline::sqlite-vec",
"inline::qdrant",
"remote::qdrant",
"remote::milvus",
"inline::milvus",
"remote::pgvector",
Expand Down
47 changes: 31 additions & 16 deletions tests/unit/providers/vector_io/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why change from tmp_path_factory?

"""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(
Copy link
Collaborator

Choose a reason for hiding this comment

The 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):
Expand Down
86 changes: 86 additions & 0 deletions tests/unit/providers/vector_io/test_qdrant.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this test actually needed?

Copy link
Contributor Author

Choose a reason for hiding this comment

The 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 QueryChunksResponse is returned without crashing.

Copy link
Contributor Author

Choose a reason for hiding this comment

The 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
Loading