unstructured-ingest 1.2.32__py3-none-any.whl

unstructured_ingest/processes/connectors/qdrant/cloud.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.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/processes/connectors/qdrant/local.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.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/processes/connectors/qdrant/qdrant.py

@@ -0,0 +1,163 @@
+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.data_types.file_data import FileData
+from unstructured_ingest.error import DestinationConnectionError, WriteError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import (
+    batch_generator,
+    flatten_dict,
+    get_enhanced_element_id,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+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/processes/connectors/qdrant/server.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.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,
+)

unstructured_ingest/processes/connectors/redisdb.py

@@ -0,0 +1,214 @@
+import json
+from contextlib import asynccontextmanager, contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Generator, Optional
+
+from pydantic import Field, Secret, model_validator
+
+from unstructured_ingest.data_types.file_data import FileData
+from unstructured_ingest.error import DestinationConnectionError, ResponseError, ValueError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Uploader,
+    UploaderConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from redis.asyncio import Redis
+
+import asyncio
+
+CONNECTOR_TYPE = "redis"
+SERVER_API_VERSION = "1"
+
+
+class RedisAccessConfig(AccessConfig):
+    uri: Optional[str] = Field(
+        default=None, description="If not anonymous, use this uri, if specified."
+    )
+    password: Optional[str] = Field(
+        default=None,
+        description="Password used to connect to database if uri is "
+        "not specified and connection is not anonymous.",
+    )
+
+
+class RedisConnectionConfig(ConnectionConfig):
+    access_config: Secret[RedisAccessConfig] = Field(
+        default=RedisAccessConfig(), validate_default=True
+    )
+    host: Optional[str] = Field(
+        default=None,
+        description="Hostname or IP address of a Redis instance to connect to "
+        "if uri is not specified.",
+    )
+    database: int = Field(default=0, description="Database index to connect to.")
+    port: Optional[int] = Field(
+        default=6379, description="Port used to connect to database if uri is not specified."
+    )
+    username: Optional[str] = Field(
+        default=None, description="Username used to connect to database if uri is not specified."
+    )
+    ssl: Optional[bool] = Field(
+        default=True,
+        description="Whether the connection should use SSL encryption if uri is not specified.",
+    )
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+
+    @model_validator(mode="after")
+    def validate_host_or_url(self) -> "RedisConnectionConfig":
+        if not self.access_config.get_secret_value().uri:
+            if not self.host:
+                raise ValueError("Please pass a hostname either directly or through uri")
+            if self.port is None:
+                raise ValueError("Since URI is not specified, port cannot be None")
+            if self.ssl is None:
+                raise ValueError("Since URI is not specified, ssl cannot be None")
+        return self
+
+    @requires_dependencies(["redis"], extras="redis")
+    @asynccontextmanager
+    async def create_async_client(self) -> AsyncGenerator["Redis", None]:
+        from redis.asyncio import Redis, from_url
+
+        access_config = self.access_config.get_secret_value()
+
+        if access_config.uri:
+            async with from_url(access_config.uri) as client:
+                yield client
+        else:
+            options = {
+                "host": self.host,
+                "port": self.port,
+                "db": self.database,
+                "ssl": self.ssl,
+                "username": self.username,
+            }
+
+            if access_config.password:
+                options["password"] = access_config.password
+            async with Redis(**options) as client:
+                yield client
+
+    @requires_dependencies(["redis"], extras="redis")
+    @contextmanager
+    def create_client(self) -> Generator["Redis", None, None]:
+        from redis import Redis, from_url
+
+        access_config = self.access_config.get_secret_value()
+
+        options = {
+            "host": self.host,
+            "port": self.port,
+            "db": self.database,
+            "ssl": self.ssl,
+            "username": self.username,
+        }
+
+        if access_config.password:
+            options["password"] = access_config.password
+
+        if access_config.uri:
+            with from_url(access_config.uri) as client:
+                yield client
+        else:
+            with Redis(**options) as client:
+                yield client
+
+
+class RedisUploaderConfig(UploaderConfig):
+    batch_size: int = Field(default=100, description="Number of records per batch")
+    key_prefix: str = Field(default="", description="Prefix for Redis keys")
+
+
+def _form_redis_pipeline_error_message(error: str) -> str:
+    """
+    Form a user-friendly error message for Redis pipeline errors.
+    The error message has `$` character at the beginning and `) of pipeline` at the end.
+    Everything between these two strings is the value an should be removed.
+    """
+    start = error.find("$")
+    end = error.find(") of pipeline")
+    if start != -1 and end != -1:
+        return error[: start + 1] + "<value>" + error[end:]
+    else:
+        return error
+
+
+@dataclass
+class RedisUploader(Uploader):
+    upload_config: RedisUploaderConfig
+    connection_config: RedisConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def is_async(self) -> bool:
+        return True
+
+    def precheck(self) -> None:
+        try:
+            with self.connection_config.create_client() as client:
+                client.ping()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    async def run_data_async(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        first_element = data[0]
+        redis_stack = await self._check_redis_stack(first_element)
+        logger.info(
+            f"writing {len(data)} objects to destination asynchronously, "
+            f"db, {self.connection_config.database}, "
+            f"at {self.connection_config.host}",
+        )
+
+        batches = list(batch_generator(data, batch_size=self.upload_config.batch_size))
+        await asyncio.gather(*[self._write_batch(batch, redis_stack) for batch in batches])
+
+    async def _write_batch(self, batch: list[dict], redis_stack: bool) -> None:
+        async with (
+            self.connection_config.create_async_client() as async_client,
+            async_client.pipeline(transaction=True) as pipe,
+        ):
+            for element in batch:
+                key_with_prefix = f"{self.upload_config.key_prefix}{element['element_id']}"
+                if redis_stack:
+                    pipe.json().set(key_with_prefix, "$", element)
+                else:
+                    pipe.set(key_with_prefix, json.dumps(element))
+            await pipe.execute()
+
+    @requires_dependencies(["redis"], extras="redis")
+    async def _check_redis_stack(self, element: dict) -> bool:
+        from redis import exceptions as redis_exceptions
+
+        redis_stack = True
+        async with (
+            self.connection_config.create_async_client() as async_client,
+            async_client.pipeline(transaction=True) as pipe,
+        ):
+            key_with_prefix = f"{self.upload_config.key_prefix}{element['element_id']}"
+            try:
+                # Redis with stack extension supports JSON type
+                await pipe.json().set(key_with_prefix, "$", element).execute()
+            except redis_exceptions.ResponseError as e:
+                message = _form_redis_pipeline_error_message(str(e))
+                if "unknown command `JSON.SET`" in message:
+                    # if this error occurs, Redis server doesn't support JSON type,
+                    # so save as string type instead
+                    await pipe.set(key_with_prefix, json.dumps(element)).execute()
+                    redis_stack = False
+                else:
+                    raise ResponseError(message) from e
+        return redis_stack
+
+
+redis_destination_entry = DestinationRegistryEntry(
+    connection_config=RedisConnectionConfig,
+    uploader=RedisUploader,
+    uploader_config=RedisUploaderConfig,
+)