unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/connector/qdrant.py

@@ -0,0 +1,144 @@
+import json
+import multiprocessing as mp
+import typing as t
+import uuid
+from dataclasses import dataclass
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError, WriteError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    ConfigSessionHandleMixin,
+    IngestDocSessionHandleMixin,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import batch_generator, flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from qdrant_client import QdrantClient
+
+
+@dataclass
+class QdrantAccessConfig(AccessConfig):
+    api_key: t.Optional[str] = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleQdrantConfig(ConfigSessionHandleMixin, BaseConnectorConfig):
+    collection_name: str
+    location: t.Optional[str] = None
+    url: t.Optional[str] = None
+    port: t.Optional[int] = 6333
+    grpc_port: t.Optional[int] = 6334
+    prefer_grpc: t.Optional[bool] = False
+    https: t.Optional[bool] = None
+    prefix: t.Optional[str] = None
+    timeout: t.Optional[float] = None
+    host: t.Optional[str] = None
+    path: t.Optional[str] = None
+    force_disable_check_same_thread: t.Optional[bool] = False
+    access_config: t.Optional[QdrantAccessConfig] = None
+
+
+@dataclass
+class QdrantWriteConfig(WriteConfig):
+    batch_size: int = 50
+    num_processes: int = 1
+
+
+@dataclass
+class QdrantDestinationConnector(IngestDocSessionHandleMixin, BaseDestinationConnector):
+    write_config: QdrantWriteConfig
+    connector_config: SimpleQdrantConfig
+    _client: t.Optional["QdrantClient"] = None
+
+    @property
+    def qdrant_client(self):
+        if self._client is None:
+            self._client = self.create_client()
+        return self._client
+
+    def initialize(self):
+        ...  # fmt: skip
+
+    @requires_dependencies(["qdrant_client"], extras="qdrant")
+    def create_client(self) -> "QdrantClient":
+        from qdrant_client import QdrantClient
+
+        client = QdrantClient(
+            location=self.connector_config.location,
+            url=self.connector_config.url,
+            port=self.connector_config.port,
+            grpc_port=self.connector_config.grpc_port,
+            prefer_grpc=self.connector_config.prefer_grpc,
+            https=self.connector_config.https,
+            api_key=(
+                self.connector_config.access_config.api_key
+                if self.connector_config.access_config
+                else None
+            ),
+            prefix=self.connector_config.prefix,
+            timeout=self.connector_config.timeout,
+            host=self.connector_config.host,
+            path=self.connector_config.path,
+            force_disable_check_same_thread=self.connector_config.force_disable_check_same_thread,
+        )
+
+        return client
+
+    @DestinationConnectionError.wrap
+    def check_connection(self):
+        self.qdrant_client.get_collections()
+
+    @DestinationConnectionError.wrap
+    @requires_dependencies(["qdrant_client"], extras="qdrant")
+    def upsert_batch(self, batch: t.List[t.Dict[str, t.Any]]):
+        from qdrant_client import models
+
+        client = self.qdrant_client
+        try:
+            points: list[models.PointStruct] = [models.PointStruct(**item) for item in batch]
+            response = client.upsert(
+                self.connector_config.collection_name, points=points, wait=True
+            )
+        except Exception as api_error:
+            raise WriteError(f"Qdrant error: {api_error}") from api_error
+        logger.debug(f"results: {response}")
+
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        logger.info(
+            f"Upserting {len(elements_dict)} elements to "
+            f"{self.connector_config.collection_name}",
+        )
+
+        qdrant_batch_size = self.write_config.batch_size
+
+        logger.info(f"using {self.write_config.num_processes} processes to upload")
+        if self.write_config.num_processes == 1:
+            for chunk in batch_generator(elements_dict, qdrant_batch_size):
+                self.upsert_batch(chunk)
+
+        else:
+            with mp.Pool(
+                processes=self.write_config.num_processes,
+            ) as pool:
+                pool.map(self.upsert_batch, list(batch_generator(elements_dict, qdrant_batch_size)))
+
+    def normalize_dict(self, element_dict: dict) -> dict:
+        return {
+            "id": str(uuid.uuid4()),
+            "vector": element_dict.pop("embeddings", {}),
+            "payload": {
+                "text": element_dict.pop("text", None),
+                "element_serialized": json.dumps(element_dict),
+                **flatten_dict(
+                    element_dict,
+                    separator="-",
+                    flatten_lists=True,
+                ),
+            },
+        }

unstructured_ingest/connector/reddit.py

@@ -0,0 +1,166 @@
+import typing as t
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    IngestDocCleanupMixin,
+    SourceConnectorCleanupMixin,
+    SourceMetadata,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from praw import Reddit
+
+
+@dataclass
+class RedditAccessConfig(AccessConfig):
+    client_secret: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+
+
+@dataclass
+class SimpleRedditConfig(BaseConnectorConfig):
+    access_config: RedditAccessConfig
+    subreddit_name: str
+    num_posts: int
+    user_agent: str
+    client_id: str
+    search_query: t.Optional[str] = None
+
+    def __post_init__(self):
+        if self.num_posts <= 0:
+            raise ValueError("The number of Reddit posts to fetch must be positive.")
+
+
+@dataclass
+class RedditIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleRedditConfig = field(repr=False)
+    post_id: str
+    registry_name: str = "reddit"
+
+    def _create_full_tmp_dir_path(self):
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+
+    @SourceConnectionNetworkError.wrap
+    @requires_dependencies(["praw"])
+    def get_post(self):
+        from praw import Reddit
+        from praw.models import Submission
+
+        reddit = Reddit(
+            client_id=self.connector_config.client_id,
+            client_secret=self.connector_config.access_config.client_secret,
+            user_agent=self.connector_config.user_agent,
+        )
+        post = Submission(reddit, self.post_id)
+        return post
+
+    def update_source_metadata(self, **kwargs):
+        post = kwargs.get("post", self.get_post())
+        if post is None:
+            self.source_metadata = SourceMetadata(
+                exists=False,
+            )
+            return
+
+        file_exists = (post.author != "[deleted]" or post.auth is not None) and (
+            post.selftext != "[deleted]" or post.selftext != "[removed]"
+        )
+
+        self.source_metadata = SourceMetadata(
+            date_created=datetime.utcfromtimestamp(post.created_utc).isoformat(),
+            source_url=post.permalink,
+            exists=file_exists,
+        )
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        """Fetches the "remote" doc and stores it locally on the filesystem."""
+        self._create_full_tmp_dir_path()
+        # Write the title plus the body, if any
+        post = self.get_post()
+        self.update_source_metadata(post=post)
+        if post is None:
+            raise ValueError(
+                f"Failed to retrieve post {self.post_id}",
+            )
+
+        text_to_write = f"# {post.title}\n{post.selftext}"
+        with open(self.filename, "w", encoding="utf8") as f:
+            f.write(text_to_write)
+
+    @property
+    def filename(self) -> Path:
+        return (Path(self.read_config.download_dir) / f"{self.post_id}.md").resolve()
+
+    @property
+    def _output_filename(self):
+        return Path(self.processor_config.output_dir) / f"{self.post_id}.json"
+
+    @property
+    def date_modified(self) -> t.Optional[str]:
+        return None
+
+    @property
+    def version(self) -> t.Optional[str]:
+        return None
+
+
+@dataclass
+class RedditSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    connector_config: SimpleRedditConfig
+    _reddit: t.Optional["Reddit"] = field(init=False, default=None)
+
+    @property
+    def reddit(self) -> "Reddit":
+        from praw import Reddit
+
+        if self._reddit is None:
+            self._reddit = Reddit(
+                client_id=self.connector_config.client_id,
+                client_secret=self.connector_config.access_config.client_secret,
+                user_agent=self.connector_config.user_agent,
+            )
+        return self._reddit
+
+    @requires_dependencies(["praw"], extras="reddit")
+    def initialize(self):
+        _ = self.reddit
+
+    def check_connection(self):
+        from praw.endpoints import API_PATH
+        from prawcore import ResponseException
+
+        try:
+            self.reddit._objectify_request(method="HEAD", params=None, path=API_PATH["me"])
+        except ResponseException as response_error:
+            logger.error(f"failed to validate connection: {response_error}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {response_error}")
+
+    def get_ingest_docs(self):
+        subreddit = self.reddit.subreddit(self.connector_config.subreddit_name)
+        if self.connector_config.search_query:
+            posts = subreddit.search(
+                self.connector_config.search_query,
+                limit=self.connector_config.num_posts,
+            )
+        else:
+            posts = subreddit.hot(limit=self.connector_config.num_posts)
+        return [
+            RedditIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                post_id=post.id,
+            )
+            for post in posts
+        ]

unstructured_ingest/connector/registry.py

@@ -0,0 +1,109 @@
+import json
+from typing import Dict, Type, cast
+
+from unstructured_ingest.connector.airtable import AirtableIngestDoc
+from unstructured_ingest.connector.astradb import AstraDBIngestDoc
+from unstructured_ingest.connector.biomed import BiomedIngestDoc
+from unstructured_ingest.connector.confluence import ConfluenceIngestDoc
+from unstructured_ingest.connector.delta_table import DeltaTableIngestDoc
+from unstructured_ingest.connector.discord import DiscordIngestDoc
+from unstructured_ingest.connector.elasticsearch import (
+    ElasticsearchIngestDoc,
+    ElasticsearchIngestDocBatch,
+)
+from unstructured_ingest.connector.fsspec.azure import AzureBlobStorageIngestDoc
+from unstructured_ingest.connector.fsspec.box import BoxIngestDoc
+from unstructured_ingest.connector.fsspec.dropbox import DropboxIngestDoc
+from unstructured_ingest.connector.fsspec.gcs import GcsIngestDoc
+from unstructured_ingest.connector.fsspec.s3 import S3IngestDoc
+from unstructured_ingest.connector.fsspec.sftp import SftpIngestDoc
+from unstructured_ingest.connector.github import GitHubIngestDoc
+from unstructured_ingest.connector.gitlab import GitLabIngestDoc
+from unstructured_ingest.connector.google_drive import GoogleDriveIngestDoc
+from unstructured_ingest.connector.hubspot import HubSpotIngestDoc
+from unstructured_ingest.connector.jira import JiraIngestDoc
+from unstructured_ingest.connector.kafka import KafkaIngestDoc
+from unstructured_ingest.connector.local import LocalIngestDoc
+from unstructured_ingest.connector.mongodb import MongoDBIngestDoc, MongoDBIngestDocBatch
+from unstructured_ingest.connector.notion.connector import (
+    NotionDatabaseIngestDoc,
+    NotionPageIngestDoc,
+)
+from unstructured_ingest.connector.onedrive import OneDriveIngestDoc
+from unstructured_ingest.connector.opensearch import OpenSearchIngestDoc, OpenSearchIngestDocBatch
+from unstructured_ingest.connector.outlook import OutlookIngestDoc
+from unstructured_ingest.connector.reddit import RedditIngestDoc
+from unstructured_ingest.connector.salesforce import SalesforceIngestDoc
+from unstructured_ingest.connector.sharepoint import SharepointIngestDoc
+from unstructured_ingest.connector.slack import SlackIngestDoc
+from unstructured_ingest.connector.wikipedia import (
+    WikipediaIngestHTMLDoc,
+    WikipediaIngestSummaryDoc,
+    WikipediaIngestTextDoc,
+)
+from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin
+from unstructured_ingest.interfaces import BaseIngestDoc
+
+INGEST_DOC_NAME_TO_CLASS: Dict[str, Type[EnhancedDataClassJsonMixin]] = {
+    "airtable": AirtableIngestDoc,
+    "astradb": AstraDBIngestDoc,
+    "azure": AzureBlobStorageIngestDoc,
+    "biomed": BiomedIngestDoc,
+    "box": BoxIngestDoc,
+    "confluence": ConfluenceIngestDoc,
+    "delta-table": DeltaTableIngestDoc,
+    "discord": DiscordIngestDoc,
+    "dropbox": DropboxIngestDoc,
+    "elasticsearch": ElasticsearchIngestDoc,
+    "elasticsearch_batch": ElasticsearchIngestDocBatch,
+    "gcs": GcsIngestDoc,
+    "github": GitHubIngestDoc,
+    "gitlab": GitLabIngestDoc,
+    "google_drive": GoogleDriveIngestDoc,
+    "hubspot": HubSpotIngestDoc,
+    "jira": JiraIngestDoc,
+    "kafka": KafkaIngestDoc,
+    "local": LocalIngestDoc,
+    "mongodb": MongoDBIngestDoc,
+    "mongodb_batch": MongoDBIngestDocBatch,
+    "notion_database": NotionDatabaseIngestDoc,
+    "notion_page": NotionPageIngestDoc,
+    "onedrive": OneDriveIngestDoc,
+    "opensearch": OpenSearchIngestDoc,
+    "opensearch_batch": OpenSearchIngestDocBatch,
+    "outlook": OutlookIngestDoc,
+    "reddit": RedditIngestDoc,
+    "s3": S3IngestDoc,
+    "salesforce": SalesforceIngestDoc,
+    "sftp": SftpIngestDoc,
+    "sharepoint": SharepointIngestDoc,
+    "slack": SlackIngestDoc,
+    "wikipedia_html": WikipediaIngestHTMLDoc,
+    "wikipedia_text": WikipediaIngestTextDoc,
+    "wikipedia_summary": WikipediaIngestSummaryDoc,
+}
+
+
+def create_ingest_doc_from_json(ingest_doc_json: str) -> BaseIngestDoc:
+    try:
+        ingest_doc_dict: dict = json.loads(ingest_doc_json)
+    except TypeError as te:
+        raise TypeError(
+            f"failed to load json string when deserializing IngestDoc: {ingest_doc_json}",
+        ) from te
+    return create_ingest_doc_from_dict(ingest_doc_dict)
+
+
+def create_ingest_doc_from_dict(ingest_doc_dict: dict) -> BaseIngestDoc:
+    ingest_doc_dict = ingest_doc_dict.copy()
+    if "registry_name" not in ingest_doc_dict:
+        raise ValueError(f"registry_name not present in ingest doc: {ingest_doc_dict}")
+    registry_name = ingest_doc_dict.pop("registry_name")
+    try:
+        ingest_doc_cls = INGEST_DOC_NAME_TO_CLASS[registry_name]
+        return cast(BaseIngestDoc, ingest_doc_cls.from_dict(ingest_doc_dict))
+    except KeyError:
+        raise ValueError(
+            f"Error: Received unknown IngestDoc name: {registry_name} while deserializing",
+            "IngestDoc.",
+        )