feat: integrate weaviate (#581)

Co-authored-by: Bogdan Banasiak <[email protected]>
Co-authored-by: Mateusz Hordyński <[email protected]>
Co-authored-by: Konrad Czarnota <[email protected]>
Co-authored-by: ds-sebastianchwilczynski <[email protected]>

Author: 5 people

.github/workflows/shared-packages.yml

@@ -144,6 +144,22 @@ jobs:
           --health-interval 10s
           --health-timeout 5s
           --health-retries 5
+      weaviate:
+        image: cr.weaviate.io/semitechnologies/weaviate:1.30.6
+        env:
+          QUERY_DEFAULTS_LIMIT: 25
+          AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED: true
+          PERSISTENCE_DATA_PATH: /var/lib/weaviate
+          ENABLE_API_BASED_MODULES: true
+          CLUSTER_HOSTNAME: node1
+        ports:
+          - 8080:8080
+          - 50051:50051
+        options: >-
+          --health-cmd "wget --no-verbose --tries=1 --spider http://localhost:8080/v1/.well-known/ready || exit 1"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
 
     steps:
       - uses: actions/checkout@v4

.libraries-whitelist.txt

@@ -10,3 +10,6 @@ mirakuru
 psycopg
 pytest-postgresql
 python-bidi
+psycopg-binary
+psycopg-pool
+griffe

docs/api_reference/core/vector-stores.md

@@ -12,4 +12,6 @@
 
 ::: ragbits.core.vector_stores.qdrant.QdrantVectorStore
 
-::: ragbits.core.vector_stores.pgvector.PgVectorStore
+::: ragbits.core.vector_stores.pgvector.PgVectorStore
+
+::: ragbits.core.vector_stores.weaviate.WeaviateVectorStore

examples/document-search/weaviate_db.py

@@ -0,0 +1,110 @@
+"""
+Ragbits Document Search Example: Weaviate
+
+This example demonstrates how to use the `DocumentSearch` class to search for documents with a more advanced setup.
+We will use the `LiteLLMEmbedder` class to embed the documents and the query, the `WeaviateVectorStore` class to store
+the embeddings.
+
+The script performs the following steps:
+
+    1. Create a list of documents.
+    2. Initialize the `LiteLLMEmbedder` class with the OpenAI `text-embedding-3-small` embedding model.
+    3. Initialize the `WeaviateVectorStore` class with a `WeaviateAsyncClient` local instance and an index name.
+    4. Initialize the `DocumentSearch` class with the embedder and the vector store.
+    5. Ingest the documents into the `DocumentSearch` instance.
+    6. List all documents in the vector store.
+    7. Search for documents using a query.
+    8. Print the list of all documents and the search results.
+
+To run the script, execute the following command:
+
+    ```bash
+    uv run examples/document-search/weaviate_db.py
+    ```
+
+Requires local Weaviate instance to be running, instructions how to set it up can be found here: https://weaviate.io/developers/weaviate/quickstart/local
+"""
+
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "ragbits-document-search",
+#     "ragbits-core[weaviate]",
+# ]
+# ///
+
+import asyncio
+
+import weaviate
+
+from ragbits.core.audit import set_trace_handlers
+from ragbits.core.embeddings.dense import LiteLLMEmbedder
+from ragbits.core.vector_stores.base import VectorStoreOptions
+from ragbits.core.vector_stores.weaviate import WeaviateVectorStore
+from ragbits.document_search import DocumentSearch, DocumentSearchOptions
+from ragbits.document_search.documents.document import DocumentMeta
+
+set_trace_handlers("cli")
+
+documents = [
+    DocumentMeta.from_literal(
+        """
+        RIP boiled water. You will be mist.
+        """
+    ),
+    DocumentMeta.from_literal(
+        """
+        Why doesn't James Bond fart in bed? Because it would blow his cover.
+        """
+    ),
+    DocumentMeta.from_literal(
+        """
+        Why programmers don't like to swim? Because they're scared of the floating points.
+        """
+    ),
+    DocumentMeta.from_literal(
+        """
+        This one is completely unrelated.
+        """
+    ),
+]
+
+
+async def main() -> None:
+    """
+    Run the example.
+    """
+    client = weaviate.use_async_with_local()
+    embedder = LiteLLMEmbedder(model_name="text-embedding-3-small")
+    vector_store = WeaviateVectorStore(
+        client=client,
+        index_name="jokes",
+        embedder=embedder,
+    )
+    document_search = DocumentSearch(
+        vector_store=vector_store,
+    )
+
+    await document_search.ingest(documents)
+
+    all_documents = await vector_store.list()
+
+    print()
+    print("All documents:")
+    print([doc.metadata["content"] for doc in all_documents])
+
+    query = "I'm boiling my water and I need a joke"
+    vector_store_options = VectorStoreOptions(
+        k=2,
+        score_threshold=0.6,
+    )
+    options = DocumentSearchOptions(vector_store_options=vector_store_options)
+    results = await document_search.search(query, options=options)
+
+    print()
+    print(f"Documents similar to: {query}")
+    print([element.text_representation for element in results])
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Integrate Weaviate vector store (#347)
 - Remove numpy dependency (#666)
 - Fix typing in LLM generate_streaming (#628)
 - Fix typing in LLM generate (#568)

packages/ragbits-core/pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "litellm>=1.55.0,<2.0.0",
     "aiohttp>=3.10.8,<4.0.0",
     "filetype>=1.2.0,<2.0.0",
+    "griffe>=1.7.3,<2.0.0"
 ]
 
 [project.urls]
@@ -94,6 +95,9 @@ hf = [
 s3 = [
     "boto3>=1.35.42,<2.0.0",
 ]
+weaviate = [
+    "weaviate-client>=4.15.4,<5.0.0",
+]
 
 [tool.uv]
 dev-dependencies = [

packages/ragbits-core/src/ragbits/core/llms/litellm.py

@@ -337,18 +337,28 @@ async def _get_litellm_response(
     ) -> ModelResponse | CustomStreamWrapper:
         entrypoint = self.router or litellm
 
+        # Prepare kwargs for the completion call
+        completion_kwargs = {
+            "messages": conversation,
+            "model": self.model_name,
+            "response_format": response_format,
+            "tools": tools,
+            "stream": stream,
+            **options.dict(),
+        }
+
+        # Only add these parameters if we're not using a router
+        # Router instances have these configured at initialization time
+        if self.router is None:
+            if self.api_base is not None:
+                completion_kwargs["base_url"] = self.api_base
+            if self.api_key is not None:
+                completion_kwargs["api_key"] = self.api_key
+            if self.api_version is not None:
+                completion_kwargs["api_version"] = self.api_version
+
         try:
-            response = await entrypoint.acompletion(
-                messages=conversation,
-                model=self.model_name,
-                base_url=self.api_base,
-                api_key=self.api_key,
-                api_version=self.api_version,
-                response_format=response_format,
-                tools=tools,
-                stream=stream,
-                **options.dict(),
-            )
+            response = await entrypoint.acompletion(**completion_kwargs)
         except litellm.openai.APIConnectionError as exc:
             raise LLMConnectionError() from exc
         except litellm.openai.APIStatusError as exc:

packages/ragbits-core/src/ragbits/core/vector_stores/base.py

@@ -27,6 +27,20 @@ class VectorStoreEntry(BaseModel):
     image_bytes: SerializableBytes | None = None
     metadata: dict = {}
 
+    @pydantic.model_validator(mode="after")
+    def validate_metadata_serializable(self) -> Self:
+        """
+        Validates that metadata is JSON serializable.
+
+        Raises:
+            ValueError: If metadata contains non-serializable values.
+        """
+        try:
+            self.model_dump_json()
+        except Exception as e:
+            raise ValueError(f"Metadata must be JSON serializable. Error: {str(e)}") from e
+        return self
+
     @pydantic.model_validator(mode="after")
     def text_or_image_required(self) -> Self:
         """