feat: add destination creation support to pinecone (#361)

* update test and implementation
changelog and version
wip

* improvements

* default vector length

* fix

* allow pipeline to pass init data downstream

* add format destination name

* version

* decorator

* tidy

---------

Co-authored-by: Roman Isecke <[email protected]>

Author: ahmetmeleq

CHANGELOG.md

@@ -1,7 +1,8 @@
-## 0.5.6-dev0
+## 0.5.6
 
 ### Enhancements
 
+* **Add support for setting up destination for Pinecone**
 * Add name formatting to Weaviate destination uploader
 
 ## 0.5.5

test/integration/connectors/test_pinecone.py

@@ -351,3 +351,37 @@ def test_pinecone_stager(
         stager=stager,
         tmp_dir=tmp_path,
     )
+
+
+@requires_env(API_KEY)
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
+def test_pinecone_create_destination(pinecone_index):
+    uploader = PineconeUploader(
+        connection_config=PineconeConnectionConfig(
+            access_config=PineconeAccessConfig(api_key=get_api_key())
+        ),
+        upload_config=PineconeUploaderConfig(),
+    )
+
+    random_id = str(uuid4()).split("-")[0]
+
+    index_name = f"test-create-destination-{random_id}"
+
+    assert not uploader.index_exists(index_name=index_name)
+
+    try:
+        uploader.create_destination(destination_name=index_name, vector_length=1536)
+    except Exception as e:
+        error_body = getattr(e, "body", None)
+        raise pytest.fail(f"failed to create destination: {e} {error_body}")
+
+    assert uploader.index_exists(index_name=index_name), "destination was not created successfully"
+
+    try:
+        pc = uploader.connection_config.get_client()
+        logger.info(f"deleting index for test create destination: {index_name}")
+        pc.delete_index(name=index_name)
+    except Exception as e:
+        raise pytest.fail(f"failed to cleanup / delete the destination: {e}")
+
+    assert not uploader.index_exists(index_name=index_name), "cleanup failed"

unstructured_ingest/__version__.py

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

unstructured_ingest/v2/interfaces/process.py

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

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, Optional, TypeVar
+from typing import Any, TypeVar
 
 from pydantic import BaseModel
 
@@ -61,6 +61,6 @@ async def run_data_async(self, data: list[dict], file_data: FileData, **kwargs:
 @dataclass
 class VectorDBUploader(Uploader, ABC):
     def create_destination(
-        self, destination_name: str = "elements", vector_length: Optional[int] = None, **kwargs: Any
+        self, vector_length: int, destination_name: str = "elements", **kwargs: Any
     ) -> bool:
         return False

unstructured_ingest/v2/pipeline/pipeline.py

@@ -126,14 +126,32 @@ def log_statuses(self):
                 for kk, vv in v.items():
                     logger.error(f"{k}: [{kk}] {vv}")
 
+    def _run_initialization(self):
+        failures = {}
+        init_kwargs = {}
+        for step in self._get_ordered_steps():
+            try:
+                step.process.init(**init_kwargs)
+                step.process.precheck()
+                # Make sure embedder dimensions available for downstream steps
+                if isinstance(step.process, Embedder):
+                    embed_dimensions = step.process.config.get_embedder().dimension
+                    init_kwargs["vector_length"] = embed_dimensions
+
+            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 initialization failure: {k}: {v}")
+            raise PipelineError("Initialization failed")
+
     def run(self):
         otel_handler = OtelHandler(otel_endpoint=self.context.otel_endpoint, log_out=logger.info)
         try:
             with otel_handler.get_tracer().start_as_current_span(
                 "ingest process", record_exception=True
             ):
-                self._run_inits()
-                self._run_prechecks()
+                self._run_initialization()
                 self._run()
         finally:
             self.log_statuses()
@@ -154,43 +172,20 @@ 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 _get_all_steps(self) -> list[PipelineStep]:
-        steps = [self.indexer_step, self.downloader_step, self.partitioner_step, self.uploader_step]
+    def _get_ordered_steps(self) -> list[PipelineStep]:
+        steps = [self.indexer_step, self.downloader_step]
+        if self.uncompress_step:
+            steps.append(self.uncompress_step)
+        steps.append(self.partitioner_step)
         if self.chunker_step:
             steps.append(self.chunker_step)
         if self.embedder_step:
             steps.append(self.embedder_step)
-        if self.uncompress_step:
-            steps.append(self.uncompress_step)
         if self.stager_step:
             steps.append(self.stager_step)
+        steps.append(self.uploader_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 self._get_all_steps():
-            try:
-                step.process.precheck()
-            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 precheck failure: {k}: {v}")
-            raise PipelineError("Precheck failed")
-
     def apply_filter(self, records: list[dict]) -> list[dict]:
         if not self.filter_step:
             return records

unstructured_ingest/v2/processes/connectors/pinecone.py

@@ -1,6 +1,7 @@
 import json
+import re
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Literal, Optional
 
 from pydantic import Field, Secret
 
@@ -13,10 +14,10 @@
     AccessConfig,
     ConnectionConfig,
     FileData,
-    Uploader,
     UploaderConfig,
     UploadStager,
     UploadStagerConfig,
+    VectorDBUploader,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
@@ -41,7 +42,7 @@ class PineconeAccessConfig(AccessConfig):
 
 
 class PineconeConnectionConfig(ConnectionConfig):
-    index_name: str = Field(description="Name of the index to connect to.")
+    index_name: Optional[str] = Field(description="Name of the index to connect to.", default=None)
     access_config: Secret[PineconeAccessConfig] = Field(
         default=PineconeAccessConfig(), validate_default=True
     )
@@ -160,18 +161,101 @@ def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
 
 
 @dataclass
-class PineconeUploader(Uploader):
+class PineconeUploader(VectorDBUploader):
     upload_config: PineconeUploaderConfig
     connection_config: PineconeConnectionConfig
     connector_type: str = CONNECTOR_TYPE
 
+    def init(self, **kwargs: Any) -> None:
+        self.create_destination(**kwargs)
+
+    def index_exists(self, index_name: Optional[str]) -> bool:
+        from pinecone.exceptions import NotFoundException
+
+        index_name = index_name or self.connection_config.index_name
+        pc = self.connection_config.get_client()
+        try:
+            pc.describe_index(index_name)
+            return True
+        except NotFoundException:
+            return False
+        except Exception as e:
+            logger.error(f"failed to check if pinecone index exists : {e}")
+            raise DestinationConnectionError(f"failed to check if pinecone index exists : {e}")
+
     def precheck(self):
         try:
-            self.connection_config.get_index()
+            # just a connection check here. not an actual index_exists check
+            self.index_exists("just-checking-our-connection")
+
+            if self.connection_config.index_name and not self.index_exists(
+                self.connection_config.index_name
+            ):
+                raise DestinationConnectionError(
+                    f"index {self.connection_config.index_name} 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 format_destination_name(self, destination_name: str) -> str:
+        # Pinecone naming requirements:
+        # can only contain lowercase letters, numbers, and hyphens
+        # must be 45 characters or less
+        formatted = re.sub(r"[^a-z0-9]", "-", destination_name.lower())
+        return formatted
+
+    def create_destination(
+        self,
+        vector_length: int,
+        destination_name: str = "elements",
+        destination_type: Literal["pod", "serverless"] = "serverless",
+        serverless_cloud: str = "aws",
+        serverless_region: str = "us-west-2",
+        pod_environment: str = "us-east1-gcp",
+        pod_type: str = "p1.x1",
+        pod_count: int = 1,
+        **kwargs: Any,
+    ) -> bool:
+        from pinecone import PodSpec, ServerlessSpec
+
+        index_name = destination_name or self.connection_config.index_name
+        index_name = self.format_destination_name(index_name)
+        self.connection_config.index_name = index_name
+
+        if not self.index_exists(index_name):
+
+            logger.info(f"creating pinecone index {index_name}")
+
+            pc = self.connection_config.get_client()
+
+            if destination_type == "serverless":
+                pc.create_index(
+                    name=destination_name,
+                    dimension=vector_length,
+                    spec=ServerlessSpec(cloud=serverless_cloud, region=serverless_region),
+                    **kwargs,
+                )
+
+                return True
+
+            elif destination_type == "pod":
+                pc.create_index(
+                    name=destination_name,
+                    dimension=vector_length,
+                    spec=PodSpec(environment=pod_environment, pod_type=pod_type, pods=pod_count),
+                    **kwargs,
+                )
+
+                return True
+
+            else:
+                raise ValueError(f"unexpected destination type: {destination_type}")
+
+        else:
+            logger.debug(f"index {index_name} already exists, skipping creation")
+            return False
+
     def pod_delete_by_record_id(self, file_data: FileData) -> None:
         logger.debug(
             f"deleting any content with metadata "
@@ -266,6 +350,10 @@ def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None
         )
         # Determine if serverless or pod based index
         pinecone_client = self.connection_config.get_client()
+
+        if not self.connection_config.index_name:
+            raise ValueError("No index name specified")
+
         index_description = pinecone_client.describe_index(name=self.connection_config.index_name)
         if "serverless" in index_description.get("spec"):
             self.serverless_delete_by_record_id(file_data=file_data)

unstructured_ingest/v2/processes/connectors/weaviate/weaviate.py

@@ -230,8 +230,8 @@ def precheck(self) -> None:
             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 init(self, **kwargs: Any) -> None:
+        self.create_destination(**kwargs)
 
     def format_destination_name(self, destination_name: str) -> str:
         # Weaviate naming requirements:

unstructured_ingest/v2/processes/embedder.py

@@ -186,7 +186,7 @@ def get_embedder(self) -> "BaseEmbeddingEncoder":
 class Embedder(BaseProcess, ABC):
     config: EmbedderConfig
 
-    def init(self, *kwargs: Any) -> None:
+    def init(self, **kwargs: Any) -> None:
         self.config.get_embedder().initialize()
 
     def run(self, elements_filepath: Path, **kwargs: Any) -> list[dict]: