unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/pinecone.py

@@ -0,0 +1,277 @@
+import json
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.data_prep import (
+    flatten_dict,
+    generator_batching_wbytes,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.constants import RECORD_ID_LABEL
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.utils import get_enhanced_element_id
+
+if TYPE_CHECKING:
+    from pinecone import Index as PineconeIndex
+    from pinecone import Pinecone
+
+
+CONNECTOR_TYPE = "pinecone"
+MAX_PAYLOAD_SIZE = 2 * 1024 * 1024  # 2MB
+MAX_POOL_THREADS = 100
+MAX_METADATA_BYTES = 40960  # 40KB https://docs.pinecone.io/reference/quotas-and-limits#hard-limits
+MAX_QUERY_RESULTS = 10000
+
+
+class PineconeAccessConfig(AccessConfig):
+    pinecone_api_key: Optional[str] = Field(
+        default=None, description="API key for Pinecone.", alias="api_key"
+    )
+
+
+class PineconeConnectionConfig(ConnectionConfig):
+    index_name: str = Field(description="Name of the index to connect to.")
+    access_config: Secret[PineconeAccessConfig] = Field(
+        default=PineconeAccessConfig(), validate_default=True
+    )
+
+    @requires_dependencies(["pinecone"], extras="pinecone")
+    def get_client(self, **index_kwargs) -> "Pinecone":
+        from pinecone import Pinecone
+
+        from unstructured_ingest import __version__ as unstructured_version
+
+        return Pinecone(
+            api_key=self.access_config.get_secret_value().pinecone_api_key,
+            source_tag=f"unstructured_ingest=={unstructured_version}",
+        )
+
+    def get_index(self, **index_kwargs) -> "PineconeIndex":
+        pc = self.get_client()
+
+        index = pc.Index(name=self.index_name, **index_kwargs)
+        logger.debug(f"connected to index: {pc.describe_index(self.index_name)}")
+        return index
+
+
+ALLOWED_FIELDS = (
+    "element_id",
+    "text",
+    "parent_id",
+    "category_depth",
+    "emphasized_text_tags",
+    "emphasized_text_contents",
+    "coordinates",
+    "last_modified",
+    "page_number",
+    "filename",
+    "is_continuation",
+    "link_urls",
+    "link_texts",
+    "text_as_html",
+)
+
+
+class PineconeUploadStagerConfig(UploadStagerConfig):
+    metadata_fields: list[str] = Field(
+        default=list(ALLOWED_FIELDS),
+        description=(
+            "which metadata from the source element to map to the payload metadata being sent to "
+            "Pinecone."
+        ),
+    )
+
+
+class PineconeUploaderConfig(UploaderConfig):
+    batch_size: Optional[int] = Field(
+        default=None,
+        description="Optional number of records per batch. Will otherwise limit by size.",
+    )
+    pool_threads: Optional[int] = Field(
+        default=1, description="Optional limit on number of threads to use for upload"
+    )
+    namespace: Optional[str] = Field(
+        default=None,
+        description="The namespace to write to. If not specified, the default namespace is used",
+    )
+    record_id_key: str = Field(
+        default=RECORD_ID_LABEL,
+        description="searchable key to find entries for the same record on previous runs",
+    )
+
+
+@dataclass
+class PineconeUploadStager(UploadStager):
+    upload_stager_config: PineconeUploadStagerConfig = field(
+        default_factory=lambda: PineconeUploadStagerConfig()
+    )
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        embeddings = element_dict.pop("embeddings", None)
+        metadata: dict[str, Any] = element_dict.pop("metadata", {})
+        data_source = metadata.pop("data_source", {})
+        coordinates = metadata.pop("coordinates", {})
+        pinecone_metadata = {}
+        for possible_meta in [element_dict, metadata, data_source, coordinates]:
+            pinecone_metadata.update(
+                {
+                    k: v
+                    for k, v in possible_meta.items()
+                    if k in self.upload_stager_config.metadata_fields
+                }
+            )
+
+        metadata = flatten_dict(
+            pinecone_metadata,
+            separator="-",
+            flatten_lists=True,
+            remove_none=True,
+        )
+        metadata_size_bytes = len(json.dumps(metadata).encode())
+        if metadata_size_bytes > MAX_METADATA_BYTES:
+            logger.info(
+                f"Metadata size is {metadata_size_bytes} bytes, which exceeds the limit of"
+                f" {MAX_METADATA_BYTES} bytes per vector. Dropping the metadata."
+            )
+            metadata = {}
+
+        metadata[RECORD_ID_LABEL] = file_data.identifier
+
+        # To support more optimal deletes, a prefix is suggested for each record:
+        # https://docs.pinecone.io/guides/data/manage-rag-documents#delete-all-records-for-a-parent-document
+        return {
+            "id": f"{file_data.identifier}#{get_enhanced_element_id(element_dict=element_dict, file_data=file_data)}",  # noqa:E501
+            "values": embeddings,
+            "metadata": metadata,
+        }
+
+
+@dataclass
+class PineconeUploader(Uploader):
+    upload_config: PineconeUploaderConfig
+    connection_config: PineconeConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def precheck(self):
+        try:
+            self.connection_config.get_index()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def pod_delete_by_record_id(self, file_data: FileData) -> None:
+        logger.debug(
+            f"deleting any content with metadata "
+            f"{self.upload_config.record_id_key}={file_data.identifier} "
+            f"from pinecone pod index"
+        )
+        index = self.connection_config.get_index(pool_threads=MAX_POOL_THREADS)
+        delete_kwargs = {
+            "filter": {self.upload_config.record_id_key: {"$eq": file_data.identifier}}
+        }
+        if namespace := self.upload_config.namespace:
+            delete_kwargs["namespace"] = namespace
+
+        resp = index.delete(**delete_kwargs)
+        logger.debug(
+            f"deleted any content with metadata "
+            f"{self.upload_config.record_id_key}={file_data.identifier} "
+            f"from pinecone index: {resp}"
+        )
+
+    def serverless_delete_by_record_id(self, file_data: FileData) -> None:
+        logger.debug(
+            f"deleting any content with metadata "
+            f"{self.upload_config.record_id_key}={file_data.identifier} "
+            f"from pinecone serverless index"
+        )
+        index = self.connection_config.get_index(pool_threads=MAX_POOL_THREADS)
+        list_kwargs = {"prefix": f"{file_data.identifier}#"}
+        deleted_ids = 0
+        if namespace := self.upload_config.namespace:
+            list_kwargs["namespace"] = namespace
+        for ids in index.list(**list_kwargs):
+            deleted_ids += len(ids)
+            delete_kwargs = {"ids": ids}
+            if namespace := self.upload_config.namespace:
+                delete_resp = delete_kwargs["namespace"] = namespace
+                # delete_resp should be an empty dict if there were no errors
+                if delete_resp:
+                    logger.error(f"failed to delete batch of ids: {delete_resp}")
+            index.delete(**delete_kwargs)
+        logger.info(
+            f"deleted {deleted_ids} records with metadata "
+            f"{self.upload_config.record_id_key}={file_data.identifier} "
+            f"from pinecone index"
+        )
+
+    @requires_dependencies(["pinecone"], extras="pinecone")
+    def upsert_batches_async(self, elements_dict: list[dict]):
+        from pinecone.exceptions import PineconeApiException
+
+        chunks = list(
+            generator_batching_wbytes(
+                iterable=elements_dict,
+                batch_size_limit_bytes=MAX_PAYLOAD_SIZE - 100,
+                max_batch_size=self.upload_config.batch_size,
+            )
+        )
+        logger.info(f"split doc with {len(elements_dict)} elements into {len(chunks)} batches")
+
+        max_pool_threads = min(len(chunks), MAX_POOL_THREADS)
+        if self.upload_config.pool_threads:
+            pool_threads = min(self.upload_config.pool_threads, max_pool_threads)
+        else:
+            pool_threads = max_pool_threads
+        index = self.connection_config.get_index(pool_threads=pool_threads)
+        with index:
+            upsert_kwargs = [{"vectors": chunk, "async_req": True} for chunk in chunks]
+            if namespace := self.upload_config.namespace:
+                for kwargs in upsert_kwargs:
+                    kwargs["namespace"] = namespace
+            async_results = [index.upsert(**kwarg) for kwarg in upsert_kwargs]
+            # Wait for and retrieve responses (this raises in case of error)
+            try:
+                results = [async_result.get() for async_result in async_results]
+            except PineconeApiException as api_error:
+                raise DestinationConnectionError(f"http error: {api_error}") from api_error
+            logger.debug(f"results: {results}")
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        logger.info(
+            f"writing a total of {len(data)} elements via"
+            f" document batches to destination"
+            f" index named {self.connection_config.index_name}"
+        )
+        # Determine if serverless or pod based index
+        pinecone_client = self.connection_config.get_client()
+        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)
+        elif "pod" in index_description.get("spec"):
+            self.pod_delete_by_record_id(file_data=file_data)
+        else:
+            raise ValueError(f"unexpected spec type in index description: {index_description}")
+        self.upsert_batches_async(elements_dict=data)
+
+
+pinecone_destination_entry = DestinationRegistryEntry(
+    connection_config=PineconeConnectionConfig,
+    uploader=PineconeUploader,
+    uploader_config=PineconeUploaderConfig,
+    upload_stager=PineconeUploadStager,
+    upload_stager_config=PineconeUploadStagerConfig,
+)

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

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    add_destination_entry,
+)
+
+from .cloud import CONNECTOR_TYPE as CLOUD_CONNECTOR_TYPE
+from .cloud import qdrant_cloud_destination_entry
+from .local import CONNECTOR_TYPE as LOCAL_CONNECTOR_TYPE
+from .local import qdrant_local_destination_entry
+from .server import CONNECTOR_TYPE as SERVER_CONNECTOR_TYPE
+from .server import qdrant_server_destination_entry
+
+add_destination_entry(destination_type=CLOUD_CONNECTOR_TYPE, entry=qdrant_cloud_destination_entry)
+add_destination_entry(destination_type=SERVER_CONNECTOR_TYPE, entry=qdrant_server_destination_entry)
+add_destination_entry(destination_type=LOCAL_CONNECTOR_TYPE, entry=qdrant_local_destination_entry)

unstructured_ingest/v2/processes/connectors/qdrant/cloud.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.qdrant.qdrant import (
+    QdrantAccessConfig,
+    QdrantConnectionConfig,
+    QdrantUploader,
+    QdrantUploaderConfig,
+    QdrantUploadStager,
+    QdrantUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "qdrant-cloud"
+
+
+class CloudQdrantAccessConfig(QdrantAccessConfig):
+    api_key: str = Field(description="Qdrant API key")
+
+
+class CloudQdrantConnectionConfig(QdrantConnectionConfig):
+    url: str = Field(default=None, description="url of Qdrant Cloud")
+    access_config: Secret[CloudQdrantAccessConfig]
+
+    def get_client_kwargs(self) -> dict:
+        return {
+            "api_key": self.access_config.get_secret_value().api_key,
+            "url": self.url,
+        }
+
+
+class CloudQdrantUploadStagerConfig(QdrantUploadStagerConfig):
+    pass
+
+
+@dataclass
+class CloudQdrantUploadStager(QdrantUploadStager):
+    upload_stager_config: CloudQdrantUploadStagerConfig
+
+
+class CloudQdrantUploaderConfig(QdrantUploaderConfig):
+    pass
+
+
+@dataclass
+class CloudQdrantUploader(QdrantUploader):
+    connection_config: CloudQdrantConnectionConfig
+    upload_config: CloudQdrantUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+qdrant_cloud_destination_entry = DestinationRegistryEntry(
+    connection_config=CloudQdrantConnectionConfig,
+    uploader=CloudQdrantUploader,
+    uploader_config=CloudQdrantUploaderConfig,
+    upload_stager=CloudQdrantUploadStager,
+    upload_stager_config=CloudQdrantUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/qdrant/local.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.qdrant.qdrant import (
+    QdrantAccessConfig,
+    QdrantConnectionConfig,
+    QdrantUploader,
+    QdrantUploaderConfig,
+    QdrantUploadStager,
+    QdrantUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "qdrant-local"
+
+
+class LocalQdrantAccessConfig(QdrantAccessConfig):
+    pass
+
+
+class LocalQdrantConnectionConfig(QdrantConnectionConfig):
+    path: str = Field(default=None, description="Persistence path for QdrantLocal.")
+    access_config: Secret[LocalQdrantAccessConfig] = Field(
+        default_factory=LocalQdrantAccessConfig, validate_default=True
+    )
+
+    def get_client_kwargs(self) -> dict:
+        return {"path": self.path}
+
+
+class LocalQdrantUploadStagerConfig(QdrantUploadStagerConfig):
+    pass
+
+
+@dataclass
+class LocalQdrantUploadStager(QdrantUploadStager):
+    upload_stager_config: LocalQdrantUploadStagerConfig
+
+
+class LocalQdrantUploaderConfig(QdrantUploaderConfig):
+    pass
+
+
+@dataclass
+class LocalQdrantUploader(QdrantUploader):
+    connection_config: LocalQdrantConnectionConfig
+    upload_config: LocalQdrantUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+qdrant_local_destination_entry = DestinationRegistryEntry(
+    connection_config=LocalQdrantConnectionConfig,
+    uploader=LocalQdrantUploader,
+    uploader_config=LocalQdrantUploaderConfig,
+    upload_stager=LocalQdrantUploadStager,
+    upload_stager_config=LocalQdrantUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/qdrant/qdrant.py

@@ -0,0 +1,160 @@
+import asyncio
+import json
+from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager, contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Generator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError, WriteError
+from unstructured_ingest.utils.data_prep import batch_generator, flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.utils import get_enhanced_element_id
+
+if TYPE_CHECKING:
+    from qdrant_client import AsyncQdrantClient, QdrantClient
+
+
+class QdrantAccessConfig(AccessConfig, ABC):
+    pass
+
+
+class QdrantConnectionConfig(ConnectionConfig, ABC):
+    access_config: Secret[QdrantAccessConfig] = Field(
+        default_factory=QdrantAccessConfig, validate_default=True, description="Access Config"
+    )
+
+    @abstractmethod
+    def get_client_kwargs(self) -> dict:
+        pass
+
+    @requires_dependencies(["qdrant_client"], extras="qdrant")
+    @asynccontextmanager
+    async def get_async_client(self) -> AsyncGenerator["AsyncQdrantClient", None]:
+        from qdrant_client import AsyncQdrantClient
+
+        client_kwargs = self.get_client_kwargs()
+        client = AsyncQdrantClient(**client_kwargs)
+        try:
+            yield client
+        finally:
+            await client.close()
+
+    @requires_dependencies(["qdrant_client"], extras="qdrant")
+    @contextmanager
+    def get_client(self) -> Generator["QdrantClient", None, None]:
+        from qdrant_client import QdrantClient
+
+        client_kwargs = self.get_client_kwargs()
+        client = QdrantClient(**client_kwargs)
+        try:
+            yield client
+        finally:
+            client.close()
+
+
+class QdrantUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class QdrantUploadStager(UploadStager, ABC):
+    upload_stager_config: QdrantUploadStagerConfig = field(
+        default_factory=lambda: QdrantUploadStagerConfig()
+    )
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        """Prepares dictionary in the format that Chroma requires"""
+        data = element_dict.copy()
+        return {
+            "id": get_enhanced_element_id(element_dict=data, file_data=file_data),
+            "vector": data.pop("embeddings", {}),
+            "payload": {
+                "text": data.pop("text", None),
+                "element_serialized": json.dumps(data),
+                **flatten_dict(
+                    data,
+                    separator="-",
+                    flatten_lists=True,
+                ),
+            },
+        }
+
+
+class QdrantUploaderConfig(UploaderConfig):
+    collection_name: str = Field(description="Name of the collection.")
+    batch_size: int = Field(default=50, description="Number of records per batch.")
+    num_processes: Optional[int] = Field(
+        default=1,
+        description="Optional limit on number of threads to use for upload.",
+        deprecated=True,
+    )
+
+
+@dataclass
+class QdrantUploader(Uploader, ABC):
+    upload_config: QdrantUploaderConfig
+    connection_config: QdrantConnectionConfig
+
+    @DestinationConnectionError.wrap
+    def precheck(self) -> None:
+        with self.connection_config.get_client() as client:
+            collections_response = client.get_collections()
+            collection_names = [c.name for c in collections_response.collections]
+            if self.upload_config.collection_name not in collection_names:
+                raise DestinationConnectionError(
+                    "collection '{}' not found: {}".format(
+                        self.upload_config.collection_name, ", ".join(collection_names)
+                    )
+                )
+
+    def is_async(self):
+        return True
+
+    async def run_data_async(
+        self,
+        data: list[dict],
+        file_data: FileData,
+        **kwargs: Any,
+    ) -> None:
+        batches = list(batch_generator(data, batch_size=self.upload_config.batch_size))
+        logger.debug(
+            "Elements split into %i batches of size %i.",
+            len(batches),
+            self.upload_config.batch_size,
+        )
+        await asyncio.gather(*[self._upsert_batch(batch) for batch in batches])
+
+    async def _upsert_batch(self, batch: list[dict]) -> None:
+        from qdrant_client import models
+
+        points: list[models.PointStruct] = [models.PointStruct(**item) for item in batch]
+        try:
+            logger.debug(
+                "Upserting %i points to the '%s' collection.",
+                len(points),
+                self.upload_config.collection_name,
+            )
+            async with self.connection_config.get_async_client() as async_client:
+                await async_client.upsert(
+                    self.upload_config.collection_name, points=points, wait=True
+                )
+        except Exception as api_error:
+            logger.error(
+                "Failed to upsert points to the collection due to the following error %s", api_error
+            )
+
+            raise WriteError(f"Qdrant error: {api_error}") from api_error
+
+        logger.debug("Successfully upsert points to the collection.")

unstructured_ingest/v2/processes/connectors/qdrant/server.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.qdrant.qdrant import (
+    QdrantAccessConfig,
+    QdrantConnectionConfig,
+    QdrantUploader,
+    QdrantUploaderConfig,
+    QdrantUploadStager,
+    QdrantUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "qdrant-server"
+
+
+class ServerQdrantAccessConfig(QdrantAccessConfig):
+    pass
+
+
+class ServerQdrantConnectionConfig(QdrantConnectionConfig):
+    url: str = Field(default=None, description="url of Qdrant server")
+    access_config: Secret[ServerQdrantAccessConfig] = Field(
+        default_factory=ServerQdrantAccessConfig, validate_default=True
+    )
+
+    def get_client_kwargs(self) -> dict:
+        return {
+            "url": self.url,
+        }
+
+
+class ServerQdrantUploadStagerConfig(QdrantUploadStagerConfig):
+    pass
+
+
+@dataclass
+class ServerQdrantUploadStager(QdrantUploadStager):
+    upload_stager_config: ServerQdrantUploadStagerConfig
+
+
+class ServerQdrantUploaderConfig(QdrantUploaderConfig):
+    pass
+
+
+@dataclass
+class ServerQdrantUploader(QdrantUploader):
+    connection_config: ServerQdrantConnectionConfig
+    upload_config: ServerQdrantUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+qdrant_server_destination_entry = DestinationRegistryEntry(
+    connection_config=ServerQdrantConnectionConfig,
+    uploader=ServerQdrantUploader,
+    uploader_config=ServerQdrantUploaderConfig,
+    upload_stager=ServerQdrantUploadStager,
+    upload_stager_config=ServerQdrantUploadStagerConfig,
+)