feat/add code to weaviate uploader to create default collection (#355)

* Add code to weaviate uploader to create default collection

* update changelog

* support functionality without having to instantiate the uploader

* don't create collection if it already exists

* change signature

* limit default indexed fields

* add unit test

* add optional init step

* always check if collection exists, create if needed

* add collection check to precheck only if it's not None

* add custom implementation of uploader for vectordbs

* use exists() rather than list_all()

* fix int test

Author: rbiseck3

CHANGELOG.md

@@ -1,8 +1,9 @@
-## 0.4.3-dev0
+## 0.4.3-dev1
 
 ### Enhancements
 
 * **Add support for allow list when downloading from raw html**
+* **Add support for setting up destination as part of uploader**
 
 ## 0.4.2
 

test/integration/connectors/weaviate/test_local.py

@@ -25,7 +25,7 @@ def wait_for_container(timeout: int = 10, interval: int = 1) -> None:
     start_time = time.time()
     while time.time() - start_time < timeout:
         try:
-            requests.get("http://localhost:8080/v1/.well-known/read")
+            requests.get("http://localhost:8080/v1/.well-known/read", timeout=1)
             return
         except Exception as e:
             print(f"Failed to validate container healthy, sleeping for {interval} seconds: {e}")
@@ -34,15 +34,20 @@ def wait_for_container(timeout: int = 10, interval: int = 1) -> None:
 
 
 @pytest.fixture
-def collection(collections_schema_config: dict) -> str:
+def weaviate_instance():
     with container_context(
         image="semitechnologies/weaviate:1.27.3",
         ports={8080: 8080, 50051: 50051},
-    ):
+    ) as ctx:
         wait_for_container()
-        with weaviate.connect_to_local() as weaviate_client:
-            weaviate_client.collections.create_from_dict(config=collections_schema_config)
-        yield COLLECTION_NAME
+        yield ctx
+
+
+@pytest.fixture
+def collection(weaviate_instance, collections_schema_config: dict) -> str:
+    with weaviate.connect_to_local() as weaviate_client:
+        weaviate_client.collections.create_from_dict(config=collections_schema_config)
+    return COLLECTION_NAME
 
 
 def get_count(client: WeaviateClient) -> int:
@@ -129,3 +134,19 @@ def test_weaviate_local_destination(upload_file: Path, collection: str, tmp_path
         file_data=file_data,
         expected_count=expected_count,
     )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
+def test_weaviate_local_create_destination(weaviate_instance):
+    uploader = LocalWeaviateUploader(
+        upload_config=LocalWeaviateUploaderConfig(),
+        connection_config=LocalWeaviateConnectionConfig(),
+    )
+    collection_name = "system_created"
+    created = uploader.create_destination(destination_name=collection_name)
+    assert created
+    with uploader.connection_config.get_client() as weaviate_client:
+        assert weaviate_client.collections.exists(name=collection_name)
+
+    created = uploader.create_destination(destination_name=collection_name)
+    assert not created

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.3-dev0"  # pragma: no cover
+__version__ = "0.4.3-dev1"  # pragma: no cover

unstructured_ingest/v2/interfaces/__init__.py

@@ -5,7 +5,7 @@
 from .process import BaseProcess
 from .processor import ProcessorConfig
 from .upload_stager import UploadStager, UploadStagerConfig
-from .uploader import UploadContent, Uploader, UploaderConfig
+from .uploader import UploadContent, Uploader, UploaderConfig, VectorDBUploader
 
 __all__ = [
     "DownloadResponse",
@@ -29,4 +29,5 @@
     "FileDataSourceMetadata",
     "BatchFileData",
     "BatchItem",
+    "VectorDBUploader",
 ]

unstructured_ingest/v2/interfaces/process.py

@@ -8,6 +8,9 @@ class BaseProcess(ABC):
     def is_async(self) -> bool:
         return False
 
+    def init(self, *kwargs: Any) -> None:
+        pass
+
     def precheck(self) -> None:
         pass
 

unstructured_ingest/v2/interfaces/uploader.py

@@ -1,7 +1,7 @@
 from abc import ABC
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import Any, Optional, TypeVar
 
 from pydantic import BaseModel
 
@@ -38,6 +38,11 @@ def is_batch(self) -> bool:
     def run_batch(self, contents: list[UploadContent], **kwargs: Any) -> None:
         raise NotImplementedError()
 
+    def create_destination(self, destination_name: str = "elements", **kwargs: Any) -> bool:
+        # Update the uploader config if needed with a new destination that gets created.
+        # Return a flag on if anything was created or not.
+        return False
+
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         data = get_data(path=path)
         self.run_data(data=data, file_data=file_data, **kwargs)
@@ -51,3 +56,11 @@ def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None
 
     async def run_data_async(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
         return self.run_data(data=data, file_data=file_data, **kwargs)
+
+
+@dataclass
+class VectorDBUploader(Uploader, ABC):
+    def create_destination(
+        self, destination_name: str = "elements", vector_length: Optional[int] = None, **kwargs: Any
+    ) -> bool:
+        return False

unstructured_ingest/v2/pipeline/pipeline.py

@@ -11,6 +11,7 @@
 from unstructured_ingest.v2.interfaces import ProcessorConfig, Uploader
 from unstructured_ingest.v2.logger import logger, make_default_logger
 from unstructured_ingest.v2.otel import OtelHandler
+from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.pipeline.steps.chunk import Chunker, ChunkStep
 from unstructured_ingest.v2.pipeline.steps.download import DownloaderT, DownloadStep
 from unstructured_ingest.v2.pipeline.steps.embed import Embedder, EmbedStep
@@ -91,10 +92,6 @@ def __post_init__(
         self.chunker_step = ChunkStep(process=chunker, context=self.context) if chunker else None
 
         self.embedder_step = EmbedStep(process=embedder, context=self.context) if embedder else None
-        # TODO: support initialize() call from each step process
-        # Potential long call to download embedder models, run before any fanout:
-        if embedder and embedder.config:
-            embedder.config.get_embedder().initialize()
 
         self.stager_step = UploadStageStep(process=stager, context=self.context) if stager else None
         self.uploader_step = UploadStep(process=uploader, context=self.context)
@@ -135,6 +132,7 @@ def run(self):
             with otel_handler.get_tracer().start_as_current_span(
                 "ingest process", record_exception=True
             ):
+                self._run_inits()
                 self._run_prechecks()
                 self._run()
         finally:
@@ -156,7 +154,7 @@ def clean_results(self, results: list[Any | list[Any]] | None) -> list[Any] | No
         final = [f for f in flat if f]
         return final or None
 
-    def _run_prechecks(self):
+    def _get_all_steps(self) -> list[PipelineStep]:
         steps = [self.indexer_step, self.downloader_step, self.partitioner_step, self.uploader_step]
         if self.chunker_step:
             steps.append(self.chunker_step)
@@ -166,8 +164,24 @@ def _run_prechecks(self):
             steps.append(self.uncompress_step)
         if self.stager_step:
             steps.append(self.stager_step)
+        return steps
+
+    def _run_inits(self):
+        failures = {}
+
+        for step in self._get_all_steps():
+            try:
+                step.process.init()
+            except Exception as e:
+                failures[step.process.__class__.__name__] = f"[{type(e).__name__}] {e}"
+        if failures:
+            for k, v in failures.items():
+                logger.error(f"Step init failure: {k}: {v}")
+            raise PipelineError("Init failed")
+
+    def _run_prechecks(self):
         failures = {}
-        for step in steps:
+        for step in self._get_all_steps():
             try:
                 step.process.precheck()
             except Exception as e:

unstructured_ingest/v2/processes/connectors/assets/__init__.py

@@ -0,0 +1,23 @@
+{
+    "properties": [
+        {
+            "dataType": [
+                "text"
+            ],
+            "indexFilterable": true,
+            "indexSearchable": true,
+            "name": "record_id",
+            "tokenization": "word"
+        },
+        {
+            "dataType": [
+                "text"
+            ],
+            "indexFilterable": true,
+            "indexSearchable": true,
+            "name": "text",
+            "tokenization": "word"
+        }
+    ],
+    "vectorizer": "none"
+}

unstructured_ingest/v2/processes/connectors/assets/weaviate_collection_config.json

@@ -3,6 +3,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from datetime import date, datetime
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from dateutil import parser
@@ -15,10 +16,10 @@
     AccessConfig,
     ConnectionConfig,
     FileData,
-    Uploader,
     UploaderConfig,
     UploadStager,
     UploadStagerConfig,
+    VectorDBUploader,
 )
 from unstructured_ingest.v2.logger import logger
 
@@ -160,7 +161,9 @@ def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
 
 
 class WeaviateUploaderConfig(UploaderConfig):
-    collection: str = Field(description="The name of the collection this object belongs to")
+    collection: Optional[str] = Field(
+        description="The name of the collection this object belongs to", default=None
+    )
     batch_size: Optional[int] = Field(default=None, description="Number of records per batch")
     requests_per_minute: Optional[int] = Field(default=None, description="Rate limit for upload")
     dynamic_batch: bool = Field(default=True, description="Whether to use dynamic batch")
@@ -205,17 +208,50 @@ def get_batch_client(self, client: "WeaviateClient") -> Generator["BatchClient",
 
 
 @dataclass
-class WeaviateUploader(Uploader, ABC):
+class WeaviateUploader(VectorDBUploader, ABC):
     upload_config: WeaviateUploaderConfig
     connection_config: WeaviateConnectionConfig
 
+    def _collection_exists(self, collection_name: Optional[str] = None):
+        collection_name = collection_name or self.upload_config.collection
+        with self.connection_config.get_client() as weaviate_client:
+            return weaviate_client.collections.exists(name=collection_name)
+
     def precheck(self) -> None:
         try:
             self.connection_config.get_client()
+            # only if collection name populated should we check that it exists
+            if self.upload_config.collection and not self._collection_exists():
+                raise DestinationConnectionError(
+                    f"collection '{self.upload_config.collection}' does not exist"
+                )
         except Exception as e:
             logger.error(f"Failed to validate connection {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
+    def init(self, *kwargs: Any) -> None:
+        self.create_destination()
+
+    def create_destination(
+        self, destination_name: str = "elements", vector_length: Optional[int] = None, **kwargs: Any
+    ) -> bool:
+        collection_name = self.upload_config.collection or destination_name
+        self.upload_config.collection = collection_name
+        connectors_dir = Path(__file__).parents[1]
+        collection_config_file = connectors_dir / "assets" / "weaviate_collection_config.json"
+        with collection_config_file.open() as f:
+            collection_config = json.load(f)
+        collection_config["class"] = collection_name
+        if not self._collection_exists():
+            logger.info(
+                f"creating default weaviate collection '{collection_name}' with default configs"
+            )
+            with self.connection_config.get_client() as weaviate_client:
+                weaviate_client.collections.create_from_dict(config=collection_config)
+                return True
+        logger.debug(f"collection with name '{collection_name}' already exists, skipping creation")
+        return False
+
     def check_for_errors(self, client: "WeaviateClient") -> None:
         failed_uploads = client.batch.failed_objects
         if failed_uploads:
@@ -253,6 +289,8 @@ def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None
             f"writing {len(data)} objects to destination "
             f"class {self.connection_config.access_config} "
         )
+        if not self.upload_config.collection:
+            raise ValueError("No collection specified")
 
         with self.connection_config.get_client() as weaviate_client:
             self.delete_by_record_id(client=weaviate_client, file_data=file_data)