feat: implement keyword and hybrid search for Weaviate provider

Author: ChristianZaccaria

llama_stack/providers/remote/vector_io/weaviate/weaviate.py

@@ -10,7 +10,7 @@
 import weaviate.classes as wvc
 from numpy.typing import NDArray
 from weaviate.classes.init import Auth
-from weaviate.classes.query import Filter
+from weaviate.classes.query import Filter, HybridFusion
 
 from llama_stack.apis.common.content_types import InterleavedContent
 from llama_stack.apis.common.errors import VectorStoreNotFoundError
@@ -26,6 +26,7 @@
     OpenAIVectorStoreMixin,
 )
 from llama_stack.providers.utils.memory.vector_store import (
+    RERANKER_TYPE_RRF,
     ChunkForDeletion,
     EmbeddingIndex,
     VectorDBWithIndex,
@@ -88,6 +89,9 @@ async def delete_chunks(self, chunks_for_deletion: list[ChunkForDeletion]) -> No
         collection.data.delete_many(where=Filter.by_property("chunk_id").contains_any(chunk_ids))
 
     async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
+        log.info(
+            f"WEAVIATE VECTOR SEARCH CALLED: embedding_shape={embedding.shape}, k={k}, threshold={score_threshold}"
+        )
         sanitized_collection_name = sanitize_collection_name(self.collection_name, weaviate_format=True)
         collection = self.client.collections.get(sanitized_collection_name)
 
@@ -109,12 +113,16 @@ async def query_vector(self, embedding: NDArray, k: int, score_threshold: float)
                 continue
 
             score = 1.0 / doc.metadata.distance if doc.metadata.distance != 0 else float("inf")
+            log.info(f"📈 Document distance: {doc.metadata.distance}, calculated score: {score}")
+
             if score < score_threshold:
                 continue
 
+            log.info(f"Document {chunk.metadata.get('document_id')} has score {score}")
             chunks.append(chunk)
             scores.append(score)
 
+        log.info(f"WEAVIATE VECTOR SEARCH RESULTS: Found {len(chunks)} chunks with scores {scores}")
         return QueryChunksResponse(chunks=chunks, scores=scores)
 
     async def delete(self, chunk_ids: list[str] | None = None) -> None:
@@ -136,7 +144,46 @@ async def query_keyword(
         k: int,
         score_threshold: float,
     ) -> QueryChunksResponse:
-        raise NotImplementedError("Keyword search is not supported in Weaviate")
+        """
+        Performs BM25-based keyword search using Weaviate's built-in full-text search.
+        Args:
+            query_string: The text query for keyword search
+            k: Limit of number of results to return
+            score_threshold: Minimum similarity score threshold
+        Returns:
+            QueryChunksResponse with combined results
+        """
+        log.info(f"WEAVIATE KEYWORD SEARCH CALLED: query='{query_string}', k={k}, threshold={score_threshold}")
+        sanitized_collection_name = sanitize_collection_name(self.collection_name, weaviate_format=True)
+        collection = self.client.collections.get(sanitized_collection_name)
+
+        # Perform BM25 keyword search on chunk_content field
+        results = collection.query.bm25(
+            query=query_string,
+            limit=k,
+            return_metadata=wvc.query.MetadataQuery(score=True),
+        )
+
+        chunks = []
+        scores = []
+        for doc in results.objects:
+            chunk_json = doc.properties["chunk_content"]
+            try:
+                chunk_dict = json.loads(chunk_json)
+                chunk = Chunk(**chunk_dict)
+            except Exception:
+                log.exception(f"Failed to parse document: {chunk_json}")
+                continue
+
+            score = doc.metadata.score if doc.metadata.score is not None else 0.0
+            if score < score_threshold:
+                continue
+
+            chunks.append(chunk)
+            scores.append(score)
+
+        log.info(f"WEAVIATE KEYWORD SEARCH RESULTS: Found {len(chunks)} chunks with scores {scores}.")
+        return QueryChunksResponse(chunks=chunks, scores=scores)
 
     async def query_hybrid(
         self,
@@ -147,7 +194,62 @@ async def query_hybrid(
         reranker_type: str,
         reranker_params: dict[str, Any] | None = None,
     ) -> QueryChunksResponse:
-        raise NotImplementedError("Hybrid search is not supported in Weaviate")
+        """
+        Hybrid search combining vector similarity and keyword search using Weaviate's native hybrid search.
+        Args:
+            embedding: The query embedding vector
+            query_string: The text query for keyword search
+            k: Limit of number of results to return
+            score_threshold: Minimum similarity score threshold
+            reranker_type: Type of reranker to use ("rrf" or "normalized")
+            reranker_params: Parameters for the reranker
+        Returns:
+            QueryChunksResponse with combined results
+        """
+        log.info(
+            f"WEAVIATE HYBRID SEARCH CALLED: query='{query_string}', embedding_shape={embedding.shape}, k={k}, threshold={score_threshold}, reranker={reranker_type}"
+        )
+        sanitized_collection_name = sanitize_collection_name(self.collection_name, weaviate_format=True)
+        collection = self.client.collections.get(sanitized_collection_name)
+
+        # Ranked (RRF) reranker fusion type
+        if reranker_type == RERANKER_TYPE_RRF:
+            rerank = HybridFusion.RANKED
+        # Relative score (Normalized) reranker fusion type
+        else:
+            rerank = HybridFusion.RELATIVE_SCORE
+
+        # Perform hybrid search using Weaviate's native hybrid search
+        results = collection.query.hybrid(
+            query=query_string,
+            alpha=0.5,  # Range <0, 1>, where 0.5 will equally favor vector and keyword search
+            vector=embedding.tolist(),
+            limit=k,
+            fusion_type=rerank,
+            return_metadata=wvc.query.MetadataQuery(score=True),
+        )
+
+        chunks = []
+        scores = []
+        for doc in results.objects:
+            chunk_json = doc.properties["chunk_content"]
+            try:
+                chunk_dict = json.loads(chunk_json)
+                chunk = Chunk(**chunk_dict)
+            except Exception:
+                log.exception(f"Failed to parse document: {chunk_json}")
+                continue
+
+            score = doc.metadata.score if doc.metadata.score is not None else 0.0
+            if score < score_threshold:
+                continue
+
+            log.info(f"Document {chunk.metadata.get('document_id')} has score {score}")
+            chunks.append(chunk)
+            scores.append(score)
+
+        log.info(f"WEAVIATE HYBRID SEARCH RESULTS: Found {len(chunks)} chunks with scores {scores}")
+        return QueryChunksResponse(chunks=chunks, scores=scores)
 
 
 class WeaviateVectorIOAdapter(

llama_stack/providers/utils/memory/vector_store.py

@@ -50,6 +50,7 @@ class ChunkForDeletion(BaseModel):
 # Constants for reranker types
 RERANKER_TYPE_RRF = "rrf"
 RERANKER_TYPE_WEIGHTED = "weighted"
+RERANKER_TYPE_NORMALIZED = "normalized"
 
 
 def parse_pdf(data: bytes) -> str:
@@ -325,6 +326,8 @@ async def query_chunks(
                 weights = ranker.get("params", {}).get("weights", [0.5, 0.5])
                 reranker_type = RERANKER_TYPE_WEIGHTED
                 reranker_params = {"alpha": weights[0] if len(weights) > 0 else 0.5}
+            elif strategy == "normalized":
+                reranker_type = RERANKER_TYPE_NORMALIZED
             else:
                 reranker_type = RERANKER_TYPE_RRF
                 k_value = ranker.get("params", {}).get("k", 60.0)

pyproject.toml

@@ -25,8 +25,8 @@ classifiers = [
 ]
 dependencies = [
     "aiohttp",
-    "fastapi>=0.115.0,<1.0",                  # server
-    "fire",                                   # for MCP in LLS client
+    "fastapi>=0.115.0,<1.0", # server
+    "fire", # for MCP in LLS client
     "httpx",
     "huggingface-hub>=0.34.0,<1.0",
     "jinja2>=3.1.6",
@@ -44,12 +44,13 @@ dependencies = [
     "tiktoken",
     "pillow",
     "h11>=0.16.0",
-    "python-multipart>=0.0.20",               # For fastapi Form
-    "uvicorn>=0.34.0",                        # server
-    "opentelemetry-sdk>=1.30.0",              # server
+    "python-multipart>=0.0.20", # For fastapi Form
+    "uvicorn>=0.34.0", # server
+    "opentelemetry-sdk>=1.30.0", # server
     "opentelemetry-exporter-otlp-proto-http>=1.30.0", # server
-    "aiosqlite>=0.21.0",                      # server - for metadata store
-    "asyncpg",                                # for metadata store
+    "aiosqlite>=0.21.0", # server - for metadata store
+    "asyncpg", # for metadata store
+    "weaviate-client>=4.16.5",
 ]
 
 [project.optional-dependencies]

tests/integration/vector_io/test_openai_vector_stores.py

@@ -22,16 +22,16 @@ def skip_if_provider_doesnt_support_openai_vector_stores(client_with_models):
     vector_io_providers = [p for p in client_with_models.providers.list() if p.api == "vector_io"]
     for p in vector_io_providers:
         if p.provider_type in [
+            "inline::chromadb",
             "inline::faiss",
-            "inline::sqlite-vec",
             "inline::milvus",
-            "inline::chromadb",
-            "remote::pgvector",
+            "inline::qdrant",
+            "inline::sqlite-vec",
             "remote::chromadb",
+            "remote::milvus",
+            "remote::pgvector",
             "remote::qdrant",
-            "inline::qdrant",
             "remote::weaviate",
-            "remote::milvus",
         ]:
             return
 
@@ -47,21 +47,23 @@ def skip_if_provider_doesnt_support_openai_vector_stores_search(client_with_mode
             "inline::milvus",
             "inline::chromadb",
             "inline::qdrant",
-            "remote::pgvector",
             "remote::chromadb",
-            "remote::weaviate",
-            "remote::qdrant",
             "remote::milvus",
+            "remote::pgvector",
+            "remote::qdrant",
+            "remote::weaviate",
         ],
         "keyword": [
+            "inline::milvus",
             "inline::sqlite-vec",
             "remote::milvus",
-            "inline::milvus",
+            "remote::weaviate",
         ],
         "hybrid": [
-            "inline::sqlite-vec",
             "inline::milvus",
+            "inline::sqlite-vec",
             "remote::milvus",
+            "remote::weaviate",
         ],
     }
     supported_providers = search_mode_support.get(search_mode, [])

tests/unit/providers/vector_io/conftest.py

@@ -23,13 +23,15 @@
 from llama_stack.providers.remote.vector_io.chroma.chroma import ChromaIndex, ChromaVectorIOAdapter, maybe_await
 from llama_stack.providers.remote.vector_io.milvus.milvus import MilvusIndex, MilvusVectorIOAdapter
 from llama_stack.providers.remote.vector_io.qdrant.qdrant import QdrantVectorIOAdapter
+from llama_stack.providers.remote.vector_io.weaviate.config import WeaviateVectorIOConfig
+from llama_stack.providers.remote.vector_io.weaviate.weaviate import WeaviateIndex, WeaviateVectorIOAdapter
 
 EMBEDDING_DIMENSION = 384
 COLLECTION_PREFIX = "test_collection"
 MILVUS_ALIAS = "test_milvus"
 
 
-@pytest.fixture(params=["milvus", "sqlite_vec", "faiss", "chroma"])
+@pytest.fixture(params=["milvus", "sqlite_vec", "faiss", "chroma", "weaviate"])
 def vector_provider(request):
     return request.param
 
@@ -333,6 +335,78 @@ async def qdrant_vec_index(qdrant_vec_db_path, embedding_dimension):
     await index.delete()
 
 
+@pytest.fixture
+def weaviate_vec_db_path():
+    return "localhost:8080"
+
+
+@pytest.fixture
+async def weaviate_vec_index(weaviate_vec_db_path, embedding_dimension):
+    import uuid
+
+    import weaviate
+
+    # Connect to local Weaviate instance
+    client = weaviate.connect_to_local(
+        host="localhost",
+        port=8080,
+    )
+
+    collection_name = f"{COLLECTION_PREFIX}_{uuid.uuid4()}"
+    index = WeaviateIndex(client=client, collection_name=collection_name)
+
+    # Create the collection for this test
+    import weaviate.classes as wvc
+    from weaviate.collections.classes.config import _CollectionConfig
+
+    from llama_stack.providers.utils.vector_io.vector_utils import sanitize_collection_name
+
+    sanitized_name = sanitize_collection_name(collection_name, weaviate_format=True)
+    collection_config = _CollectionConfig(
+        name=sanitized_name,
+        vectorizer_config=wvc.config.Configure.Vectorizer.none(),
+        properties=[
+            wvc.config.Property(
+                name="chunk_content",
+                data_type=wvc.config.DataType.TEXT,
+            ),
+        ],
+    )
+    if not client.collections.exists(sanitized_name):
+        client.collections.create_from_config(collection_config)
+
+    yield index
+    await index.delete()
+    client.close()
+
+
+@pytest.fixture
+async def weaviate_vec_adapter(weaviate_vec_db_path, mock_inference_api, embedding_dimension):
+    config = WeaviateVectorIOConfig(
+        weaviate_cluster_url=weaviate_vec_db_path,
+        weaviate_api_key=None,
+        kvstore=SqliteKVStoreConfig(),
+    )
+    adapter = WeaviateVectorIOAdapter(
+        config=config,
+        inference_api=mock_inference_api,
+        files_api=None,
+    )
+    collection_id = f"weaviate_test_collection_{random.randint(1, 1_000_000)}"
+    await adapter.initialize()
+    await adapter.register_vector_db(
+        VectorDB(
+            identifier=collection_id,
+            provider_id="test_provider",
+            embedding_model="test_model",
+            embedding_dimension=embedding_dimension,
+        )
+    )
+    adapter.test_collection_id = collection_id
+    yield adapter
+    await adapter.shutdown()
+
+
 @pytest.fixture
 def vector_io_adapter(vector_provider, request):
     """Returns the appropriate vector IO adapter based on the provider parameter."""
@@ -342,6 +416,7 @@ def vector_io_adapter(vector_provider, request):
         "sqlite_vec": "sqlite_vec_adapter",
         "chroma": "chroma_vec_adapter",
         "qdrant": "qdrant_vec_adapter",
+        "weaviate": "weaviate_vec_adapter",
     }
     return request.getfixturevalue(vector_provider_dict[vector_provider])
 

tests/unit/providers/vector_io/remote/test_milvus.py

@@ -23,13 +23,13 @@
 with patch.dict("sys.modules", {"pymilvus": pymilvus_mock}):
     from llama_stack.providers.remote.vector_io.milvus.milvus import MilvusIndex
 
-# This test is a unit test for the MilvusVectorIOAdapter class. This should only contain
+# This test is a unit test for the MilvusIndex class. This should only contain
 # tests which are specific to this class. More general (API-level) tests should be placed in
 # tests/integration/vector_io/
 #
 # How to run this test:
 #
-# pytest tests/unit/providers/vector_io/test_milvus.py \
+# pytest tests/unit/providers/vector_io/remote/test_milvus.py \
 # -v -s --tb=short --disable-warnings --asyncio-mode=auto
 
 MILVUS_PROVIDER = "milvus"
@@ -106,6 +106,7 @@ async def test_add_chunks(milvus_index, sample_chunks, sample_embeddings, mock_m
 
     # Verify the insert call had the right number of chunks
     insert_call = mock_milvus_client.insert.call_args
+    print(insert_call[1])
     assert len(insert_call[1]["data"]) == len(sample_chunks)
 
 
@@ -324,3 +325,6 @@ async def test_query_hybrid_search_default_rrf(
     call_args = mock_milvus_client.hybrid_search.call_args
     ranker = call_args[1]["ranker"]
     assert ranker is not None
+
+
+# TODO: Write tests for the MilvusVectorIOAdapter class.