feat(weaviate): add tenant support in write_documents and async batch write

Author: Akash504-ai

integrations/weaviate/src/haystack_integrations/document_stores/weaviate/document_store.py

@@ -177,10 +177,11 @@ def _clean_connection_settings(self) -> None:
         _class_name = self._collection_settings.get("class", "Default")
         _class_name = _class_name[0].upper() + _class_name[1:]
         self._collection_settings["class"] = _class_name
+
         # Set the properties if they're not set
         self._collection_settings["properties"] = self._collection_settings.get(
             "properties", DOCUMENT_COLLECTION_PROPERTIES
-        )
+        )   
 
     @property
     def client(self) -> weaviate.WeaviateClient:
@@ -937,11 +938,13 @@ def _handle_failed_objects(failed_objects: list[ErrorObject]) -> NoReturn:
     def _batch_write(self, documents: list[Document], tenant: str | None = None) -> int:
         """
         Writes document to Weaviate in batches.
-
-        Documents with the same id will be overwritten.
-        Raises in case of errors.
         """
 
+        # Handle tenant at collection level (NOT via kwargs)
+        collection = self.collection
+        if tenant is not None:
+            collection = collection.with_tenant(tenant)
+
         with self.client.batch.dynamic() as batch:
             for doc in documents:
                 if not isinstance(doc, Document):
@@ -950,27 +953,28 @@ def _batch_write(self, documents: list[Document], tenant: str | None = None) ->
 
                 batch.add_object(
                     properties=WeaviateDocumentStore._to_data_object(doc),
-                    collection=self.collection.name,
+                    collection=collection.name,  
                     uuid=generate_uuid5(doc.id),
                     vector=doc.embedding,
-                    tenant=tenant
                 )
+
         if failed_objects := self.client.batch.failed_objects:
             self._handle_failed_objects(failed_objects)
 
-        # If the document already exists we get no status message back from Weaviate.
-        # So we assume that all Documents were written.
         return len(documents)
 
     async def _batch_write_async(self, documents: list[Document], tenant: str | None = None) -> int:
         """
         Asynchronously writes document to Weaviate in batches.
-
-        Documents with the same id will be overwritten.
-        Raises in case of errors.
         """
+
         client = await self.async_client
 
+        # Handle tenant properly
+        collection = await self.async_collection
+        if tenant is not None:
+            collection = collection.with_tenant(tenant)
+
         async with client.batch.stream() as batch:
             for doc in documents:
                 if not isinstance(doc, Document):
@@ -979,17 +983,14 @@ async def _batch_write_async(self, documents: list[Document], tenant: str | None
 
                 await batch.add_object(
                     properties=WeaviateDocumentStore._to_data_object(doc),
-                    collection=(await self.async_collection).name,
+                    collection=collection.name,   
                     uuid=generate_uuid5(doc.id),
                     vector=doc.embedding,
-                    tenant=tenant
                 )
 
         if failed_objects := client.batch.failed_objects:
             self._handle_failed_objects(failed_objects)
 
-        # If the document already exists we get no status message back from Weaviate.
-        # So we assume that all Documents were written.
         return len(documents)
 
     def _write(self, documents: list[Document], policy: DuplicatePolicy, tenant: str | None = None) -> int:
@@ -1260,25 +1261,23 @@ async def delete_all_documents_async(self, *, recreate_index: bool = False, batc
                 )
 
     def delete_by_filter(self, filters: dict[str, Any]) -> int:
-        """
-        Deletes all documents that match the provided filters.
-
-        :param filters: The filters to apply to select documents for deletion.
-            For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
-        :returns: The number of documents deleted.
-        """
         validate_filters(filters)
 
         try:
+            collection = self.collection  # ✅ FIX
+
             weaviate_filter = convert_filters(filters)
-            result = self.collection.data.delete_many(where=weaviate_filter)
+            result = collection.data.delete_many(where=weaviate_filter)
             deleted_count = result.successful
+
             logger.info(
                 "Deleted {n_docs} documents from collection '{collection}' using filters.",
                 n_docs=deleted_count,
-                collection=self.collection.name,
+                collection=collection.name,
             )
+
             return deleted_count
+
         except weaviate.exceptions.WeaviateQueryError as e:
             msg = f"Failed to delete documents by filter in Weaviate. Error: {e.message}"
             raise DocumentStoreError(msg) from e
@@ -1315,54 +1314,42 @@ async def delete_by_filter_async(self, filters: dict[str, Any]) -> int:
             raise DocumentStoreError(msg) from e
 
     def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int:
-        """
-        Updates the metadata of all documents that match the provided filters.
-
-        :param filters: The filters to apply to select documents for updating.
-            For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
-        :param meta: The metadata fields to update. These will be merged with existing metadata.
-        :returns: The number of documents updated.
-        """
         validate_filters(filters)
 
         if not isinstance(meta, dict):
             msg = "Meta must be a dictionary"
             raise ValueError(msg)
 
         try:
+            collection = self.collection  # ✅ FIX
+
             matching_objects = self._query_with_filters(filters)
             if not matching_objects:
                 return 0
 
-            # Update each object with the new metadata
-            # Since metadata is stored flattened in Weaviate properties, we update properties directly
             updated_count = 0
             failed_updates = []
 
             for obj in matching_objects:
                 try:
-                    # Get current properties
                     current_properties = obj.properties.copy() if obj.properties else {}
 
-                    # Update with new metadata values
-                    # Note: metadata fields are stored directly in properties (flattened)
                     for key, value in meta.items():
                         current_properties[key] = value
 
-                    # Update the object, preserving the vector
-                    # Get the vector from the object to preserve it during replace
                     vector: VECTORS | None = None
                     if isinstance(obj.vector, (list, dict)):
                         vector = obj.vector
 
-                    self.collection.data.replace(
+                    collection.data.replace(   # ✅ FIX
                         uuid=obj.uuid,
                         properties=current_properties,
                         vector=vector,
                     )
+
                     updated_count += 1
+
                 except Exception as e:
-                    # Collect failed updates but continue with others
                     obj_properties = obj.properties or {}
                     id_ = obj_properties.get("_original_id", obj.uuid)
                     failed_updates.append((id_, str(e)))
@@ -1376,9 +1363,11 @@ def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int
             logger.info(
                 "Updated {n_docs} documents in collection '{collection}' using filters.",
                 n_docs=updated_count,
-                collection=self.collection.name,
+                collection=collection.name,
             )
+
             return updated_count
+
         except weaviate.exceptions.WeaviateQueryError as e:
             msg = f"Failed to update documents by filter in Weaviate. Error: {e.message}"
             raise DocumentStoreError(msg) from e
@@ -1618,4 +1607,4 @@ async def _hybrid_retrieval_async(
             return_metadata=["score"],
         )
 
-        return [WeaviateDocumentStore._to_document(doc) for doc in result.objects]
+        return [WeaviateDocumentStore._to_document(doc) for doc in result.objects]

integrations/weaviate/tests/test_document_store.py

@@ -8,6 +8,7 @@
 from collections.abc import Generator
 from unittest.mock import MagicMock, patch
 
+import platform
 import pytest
 from dateutil import parser
 from haystack.dataclasses.byte_stream import ByteStream
@@ -833,6 +834,7 @@ def test_connect_to_local(self):
         )
         assert document_store.client
 
+    @pytest.mark.skipif(platform.system() == "Windows", reason="EmbeddedDB not supported on Windows")
     def test_connect_to_embedded(self):
         document_store = WeaviateDocumentStore(embedded_options=EmbeddedOptions())
         assert document_store.client
@@ -1142,4 +1144,4 @@ def test_count_unique_metadata_by_filter_all_documents(document_store):
         )
         assert counts["category"] == 3
         assert counts["status"] == 2
-        assert counts["priority"] == 3
+        assert counts["priority"] == 3