unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/chroma.py

@@ -0,0 +1,190 @@
+from dataclasses import dataclass, field
+from datetime import date, datetime
+from typing import TYPE_CHECKING, Annotated, Any, Optional
+
+from dateutil import parser
+from pydantic import Field, Secret
+from pydantic.functional_validators import BeforeValidator
+
+from unstructured_ingest.error import DestinationConnectionError
+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.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.utils import get_enhanced_element_id
+
+from .utils import conform_string_to_dict
+
+if TYPE_CHECKING:
+    from chromadb import Client
+
+CONNECTOR_TYPE = "chroma"
+
+
+class ChromaAccessConfig(AccessConfig):
+    settings: Optional[Annotated[dict, BeforeValidator(conform_string_to_dict)]] = Field(
+        default=None, description="A dictionary of settings to communicate with the chroma server."
+    )
+    headers: Optional[Annotated[dict, BeforeValidator(conform_string_to_dict)]] = Field(
+        default=None, description="A dictionary of headers to send to the Chroma server."
+    )
+
+
+class ChromaConnectionConfig(ConnectionConfig):
+    access_config: Secret[ChromaAccessConfig] = Field(
+        default=ChromaAccessConfig(), validate_default=True
+    )
+    path: Optional[str] = Field(
+        default=None, description="Location where Chroma is persisted, if not connecting via http."
+    )
+    tenant: Optional[str] = Field(
+        default="default_tenant", description="The tenant to use for this client."
+    )
+    database: Optional[str] = Field(
+        default="default_database", description="The database to use for this client."
+    )
+    host: Optional[str] = Field(default=None, description="The hostname of the Chroma server.")
+    port: Optional[int] = Field(default=None, description="The port of the Chroma server.")
+    ssl: bool = Field(
+        default=False, description="Whether to use SSL to connect to the Chroma server."
+    )
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+
+    @requires_dependencies(["chromadb"], extras="chroma")
+    def get_client(self) -> "Client":
+        import chromadb
+
+        access_config = self.access_config.get_secret_value()
+        if path := self.path:
+            return chromadb.PersistentClient(
+                path=path,
+                settings=access_config.settings,
+                tenant=self.tenant,
+                database=self.database,
+            )
+
+        elif (host := self.host) and (port := self.port):
+            return chromadb.HttpClient(
+                host=host,
+                port=str(port),
+                ssl=self.ssl,
+                headers=access_config.headers,
+                settings=access_config.settings,
+                tenant=self.tenant,
+                database=self.database,
+            )
+        else:
+            raise ValueError("Chroma connector requires either path or host and port to be set.")
+
+
+class ChromaUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class ChromaUploadStager(UploadStager):
+    upload_stager_config: ChromaUploadStagerConfig = field(
+        default_factory=lambda: ChromaUploadStagerConfig()
+    )
+
+    @staticmethod
+    def parse_date_string(date_string: str) -> date:
+        try:
+            timestamp = float(date_string)
+            return datetime.fromtimestamp(timestamp)
+        except Exception as e:
+            logger.debug(f"date {date_string} string not a timestamp: {e}")
+        return parser.parse(date_string)
+
+    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),
+            "embedding": data.pop("embeddings", None),
+            "document": data.pop("text", None),
+            "metadata": flatten_dict(data, separator="-", flatten_lists=True, remove_none=True),
+        }
+
+
+class ChromaUploaderConfig(UploaderConfig):
+    collection_name: str = Field(description="The name of the Chroma collection to write into.")
+    batch_size: int = Field(default=100, description="Number of records per batch")
+
+
+@dataclass
+class ChromaUploader(Uploader):
+    connector_type: str = CONNECTOR_TYPE
+    upload_config: ChromaUploaderConfig
+    connection_config: ChromaConnectionConfig
+
+    def precheck(self) -> None:
+        try:
+            self.connection_config.get_client()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    @DestinationConnectionError.wrap
+    def upsert_batch(self, collection, batch):
+        try:
+            # Chroma wants lists even if there is only one element
+            # Upserting to prevent duplicates
+            collection.upsert(
+                ids=batch["ids"],
+                documents=batch["documents"],
+                embeddings=batch["embeddings"],
+                metadatas=batch["metadatas"],
+            )
+        except Exception as e:
+            raise ValueError(f"chroma error: {e}") from e
+
+    @staticmethod
+    def prepare_chroma_list(chunk: tuple[dict[str, Any]]) -> dict[str, list[Any]]:
+        """Helper function to break a tuple of dicts into list of parallel lists for ChromaDb.
+        ({'id':1}, {'id':2}, {'id':3}) -> {'ids':[1,2,3]}"""
+        chroma_dict = {}
+        chroma_dict["ids"] = [x.get("id") for x in chunk]
+        chroma_dict["documents"] = [x.get("document") for x in chunk]
+        chroma_dict["embeddings"] = [x.get("embedding") for x in chunk]
+        chroma_dict["metadatas"] = [x.get("metadata") for x in chunk]
+        # Make sure all lists are of the same length
+        assert (
+            len(chroma_dict["ids"])
+            == len(chroma_dict["documents"])
+            == len(chroma_dict["embeddings"])
+            == len(chroma_dict["metadatas"])
+        )
+        return chroma_dict
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        logger.info(
+            f"writing {len(data)} objects to destination "
+            f"collection {self.upload_config.collection_name} "
+            f"at {self.connection_config.host}",
+        )
+        client = self.connection_config.get_client()
+
+        collection = client.get_or_create_collection(name=self.upload_config.collection_name)
+        for chunk in batch_generator(data, self.upload_config.batch_size):
+            self.upsert_batch(collection, self.prepare_chroma_list(chunk))
+
+
+chroma_destination_entry = DestinationRegistryEntry(
+    connection_config=ChromaConnectionConfig,
+    uploader=ChromaUploader,
+    uploader_config=ChromaUploaderConfig,
+    upload_stager=ChromaUploadStager,
+    upload_stager_config=ChromaUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/confluence.py

@@ -0,0 +1,207 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Generator, List, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import SourceConnectionError
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    SourceRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from atlassian import Confluence
+
+CONNECTOR_TYPE = "confluence"
+
+
+class ConfluenceAccessConfig(AccessConfig):
+    api_token: Optional[str] = Field(description="Confluence API token", default=None)
+    access_token: Optional[str] = Field(
+        description="Confluence Personal Access Token", default=None
+    )
+
+
+class ConfluenceConnectionConfig(ConnectionConfig):
+    url: str = Field(description="URL of the Confluence instance")
+    user_email: Optional[str] = Field(description="User email for authentication", default=None)
+    access_config: Secret[ConfluenceAccessConfig] = Field(
+        description="Access configuration for Confluence"
+    )
+
+    def model_post_init(self, __context):
+        access_configs = self.access_config.get_secret_value()
+        basic_auth = self.user_email and access_configs.api_token
+        pat_auth = access_configs.access_token
+        if basic_auth and pat_auth:
+            raise ValueError("both forms of auth provided, only one allowed")
+        if not (basic_auth or pat_auth):
+            raise ValueError("neither forms of auth provided")
+
+    @requires_dependencies(["atlassian"], extras="confluence")
+    def get_client(self) -> "Confluence":
+        from atlassian import Confluence
+
+        access_configs = self.access_config.get_secret_value()
+        return Confluence(
+            url=self.url,
+            username=self.user_email,
+            password=access_configs.api_token,
+        )
+
+
+class ConfluenceIndexerConfig(IndexerConfig):
+    max_num_of_spaces: int = Field(500, description="Maximum number of spaces to index")
+    max_num_of_docs_from_each_space: int = Field(
+        100, description="Maximum number of documents to fetch from each space"
+    )
+    spaces: Optional[List[str]] = Field(None, description="List of specific space keys to index")
+
+
+@dataclass
+class ConfluenceIndexer(Indexer):
+    connection_config: ConfluenceConnectionConfig
+    index_config: ConfluenceIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def precheck(self) -> bool:
+        try:
+
+            # Attempt to retrieve a list of spaces with limit=1.
+            # This should only succeed if all creds are valid
+            client = self.connection_config.get_client()
+            client.get_all_spaces(limit=1)
+            logger.info("Connection to Confluence successful.")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to connect to Confluence: {e}", exc_info=True)
+            raise SourceConnectionError(f"Failed to connect to Confluence: {e}")
+
+    def _get_space_ids(self) -> List[str]:
+        spaces = self.index_config.spaces
+        if spaces:
+            return spaces
+        else:
+            client = self.connection_config.get_client()
+            all_spaces = client.get_all_spaces(limit=self.index_config.max_num_of_spaces)
+            space_ids = [space["key"] for space in all_spaces["results"]]
+            return space_ids
+
+    def _get_docs_ids_within_one_space(self, space_id: str) -> List[dict]:
+        client = self.connection_config.get_client()
+        pages = client.get_all_pages_from_space(
+            space=space_id,
+            start=0,
+            limit=self.index_config.max_num_of_docs_from_each_space,
+            expand=None,
+            content_type="page",
+            status=None,
+        )
+        doc_ids = [{"space_id": space_id, "doc_id": page["id"]} for page in pages]
+        return doc_ids
+
+    def run(self) -> Generator[FileData, None, None]:
+        from time import time
+
+        space_ids = self._get_space_ids()
+        for space_id in space_ids:
+            doc_ids = self._get_docs_ids_within_one_space(space_id)
+            for doc in doc_ids:
+                doc_id = doc["doc_id"]
+                # Build metadata
+                metadata = FileDataSourceMetadata(
+                    date_processed=str(time()),
+                    url=f"{self.connection_config.url}/pages/{doc_id}",
+                    record_locator={
+                        "space_id": space_id,
+                        "document_id": doc_id,
+                    },
+                )
+                additional_metadata = {
+                    "space_id": space_id,
+                    "document_id": doc_id,
+                }
+
+                # Construct relative path and filename
+                filename = f"{doc_id}.html"
+                relative_path = str(Path(space_id) / filename)
+
+                source_identifiers = SourceIdentifiers(
+                    filename=filename,
+                    fullpath=relative_path,
+                    rel_path=relative_path,
+                )
+
+                file_data = FileData(
+                    identifier=doc_id,
+                    connector_type=self.connector_type,
+                    metadata=metadata,
+                    additional_metadata=additional_metadata,
+                    source_identifiers=source_identifiers,
+                )
+                yield file_data
+
+
+class ConfluenceDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class ConfluenceDownloader(Downloader):
+    connection_config: ConfluenceConnectionConfig
+    download_config: ConfluenceDownloaderConfig = field(default_factory=ConfluenceDownloaderConfig)
+    connector_type: str = CONNECTOR_TYPE
+
+    def run(self, file_data: FileData, **kwargs) -> DownloadResponse:
+        doc_id = file_data.identifier
+        try:
+            client = self.connection_config.get_client()
+            page = client.get_page_by_id(
+                page_id=doc_id,
+                expand="history.lastUpdated,version,body.view",
+            )
+        except Exception as e:
+            logger.error(f"Failed to retrieve page with ID {doc_id}: {e}", exc_info=True)
+            raise SourceConnectionError(f"Failed to retrieve page with ID {doc_id}: {e}")
+
+        if not page:
+            raise ValueError(f"Page with ID {doc_id} does not exist.")
+
+        content = page["body"]["view"]["value"]
+
+        filepath = file_data.source_identifiers.relative_path
+        download_path = Path(self.download_dir) / filepath
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(download_path, "w", encoding="utf8") as f:
+            f.write(content)
+
+        # Update file_data with metadata
+        file_data.metadata.date_created = page["history"]["createdDate"]
+        file_data.metadata.date_modified = page["version"]["when"]
+        file_data.metadata.version = str(page["version"]["number"])
+        file_data.display_name = page["title"]
+
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+
+confluence_source_entry = SourceRegistryEntry(
+    connection_config=ConfluenceConnectionConfig,
+    indexer_config=ConfluenceIndexerConfig,
+    indexer=ConfluenceIndexer,
+    downloader_config=ConfluenceDownloaderConfig,
+    downloader=ConfluenceDownloader,
+)