unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/sql.py

@@ -0,0 +1,269 @@
+import enum
+import json
+import uuid
+from dataclasses import dataclass, field
+from datetime import date, datetime
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import numpy as np
+import pandas as pd
+from dateutil import parser
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+
+CONNECTOR_TYPE = "sql"
+ELEMENTS_TABLE_NAME = "elements"
+
+
+@dataclass
+class SQLAccessConfig(AccessConfig):
+    username: Optional[str] = None
+    password: Optional[str] = None
+
+
+class DatabaseType(str, enum.Enum):
+    SQLITE = "sqlite"
+    POSTGRESQL = "postgresql"
+
+
+@dataclass
+class SimpleSqlConfig(ConnectionConfig):
+    db_type: DatabaseType = (
+        # required default value here because of parent class
+        DatabaseType.SQLITE
+    )
+    database: Optional[str] = None
+    host: Optional[str] = None
+    port: Optional[int] = 5432
+    access_config: Optional[SQLAccessConfig] = enhanced_field(default=None, sensitive=True)
+    connector_type: str = CONNECTOR_TYPE
+
+    def __post_init__(self):
+        if (self.db_type == DatabaseType.SQLITE) and (self.database is None):
+            raise ValueError(
+                "A sqlite connection requires a path to a *.db file "
+                "through the `database` argument"
+            )
+
+
+@dataclass
+class SQLUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+_COLUMNS = (
+    "id",
+    "element_id",
+    "text",
+    "embeddings",
+    "type",
+    "system",
+    "layout_width",
+    "layout_height",
+    "points",
+    "url",
+    "version",
+    "date_created",
+    "date_modified",
+    "date_processed",
+    "permissions_data",
+    "record_locator",
+    "category_depth",
+    "parent_id",
+    "attached_filename",
+    "filetype",
+    "last_modified",
+    "file_directory",
+    "filename",
+    "languages",
+    "page_number",
+    "links",
+    "page_name",
+    "link_urls",
+    "link_texts",
+    "sent_from",
+    "sent_to",
+    "subject",
+    "section",
+    "header_footer_type",
+    "emphasized_text_contents",
+    "emphasized_text_tags",
+    "text_as_html",
+    "regex_metadata",
+    "detection_class_prob",
+)
+
+_DATE_COLUMNS = ("date_created", "date_modified", "date_processed", "last_modified")
+
+
+def parse_date_string(date_value: Union[str, int]) -> date:
+    try:
+        timestamp = float(date_value) / 1000 if isinstance(date_value, int) else float(date_value)
+        return datetime.fromtimestamp(timestamp)
+    except Exception as e:
+        logger.debug(f"date {date_value} string not a timestamp: {e}")
+    return parser.parse(date_value)
+
+
+@dataclass
+class SQLUploadStager(UploadStager):
+    upload_stager_config: SQLUploadStagerConfig = field(
+        default_factory=lambda: SQLUploadStagerConfig()
+    )
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        with open(elements_filepath) as elements_file:
+            elements_contents = json.load(elements_file)
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        output = []
+        for data in elements_contents:
+            metadata: dict[str, Any] = data.pop("metadata", {})
+            data_source = metadata.pop("data_source", {})
+            coordinates = metadata.pop("coordinates", {})
+
+            data.update(metadata)
+            data.update(data_source)
+            data.update(coordinates)
+
+            data["id"] = str(uuid.uuid4())
+
+            # remove extraneous, not supported columns
+            [data.pop(column) for column in data if column not in _COLUMNS]
+
+            output.append(data)
+
+        df = pd.DataFrame.from_dict(output)
+        for column in filter(lambda x: x in df.columns, _DATE_COLUMNS):
+            df[column] = df[column].apply(parse_date_string)
+        for column in filter(
+            lambda x: x in df.columns,
+            ("permissions_data", "record_locator", "points", "links"),
+        ):
+            df[column] = df[column].apply(
+                lambda x: json.dumps(x) if isinstance(x, (list, dict)) else None
+            )
+        for column in filter(
+            lambda x: x in df.columns,
+            ("version", "page_number", "regex_metadata"),
+        ):
+            df[column] = df[column].apply(str)
+
+        with output_path.open("w") as output_file:
+            df.to_json(output_file, orient="records", lines=True)
+        return output_path
+
+
+@dataclass
+class SQLUploaderConfig(UploaderConfig):
+    batch_size: int = 50
+
+
+@dataclass
+class SQLUploader(Uploader):
+    connector_type: str = CONNECTOR_TYPE
+    upload_config: SQLUploaderConfig
+    connection_config: SimpleSqlConfig
+
+    @property
+    def connection(self):
+        if self.connection_config.db_type == DatabaseType.POSTGRESQL:
+            return self._make_psycopg_connection
+        elif self.connection_config.db_type == DatabaseType.SQLITE:
+            return self._make_sqlite_connection
+        raise ValueError(f"Unsupported database {self.connection_config.db_type} connection.")
+
+    def _make_sqlite_connection(self):
+        from sqlite3 import connect
+
+        return connect(database=self.connection_config.database)
+
+    @requires_dependencies(["psycopg2"], extras="postgres")
+    def _make_psycopg_connection(self):
+        from psycopg2 import connect
+
+        return connect(
+            user=self.connection_config.access_config.username,
+            password=self.connection_config.access_config.password,
+            dbname=self.connection_config.database,
+            host=self.connection_config.host,
+            port=self.connection_config.port,
+        )
+
+    def prepare_data(
+        self, columns: list[str], data: tuple[tuple[Any, ...], ...]
+    ) -> list[tuple[Any, ...]]:
+        output = []
+        for row in data:
+            parsed = []
+            for column_name, value in zip(columns, row):
+                if self.connection_config.db_type == DatabaseType.SQLITE and isinstance(
+                    value, (list, dict)
+                ):
+                    value = json.dumps(value)
+                if column_name in _DATE_COLUMNS:
+                    if value is None:
+                        parsed.append(None)
+                    else:
+                        parsed.append(parse_date_string(value))
+                else:
+                    parsed.append(value)
+            output.append(tuple(parsed))
+        return output
+
+    def upload_contents(self, content: UploadContent) -> None:
+        df = pd.read_json(content.path, orient="records", lines=True)
+        logger.debug(f"uploading {len(df)} entries to {self.connection_config.database} ")
+        df.replace({np.nan: None}, inplace=True)
+
+        columns = tuple(df.columns)
+        stmt = f"INSERT INTO {ELEMENTS_TABLE_NAME} ({','.join(columns)}) \
+                VALUES({','.join(['?' if self.connection_config.db_type==DatabaseType.SQLITE else '%s' for x in columns])})"  # noqa E501
+
+        for rows in pd.read_json(
+            content.path, orient="records", lines=True, chunksize=self.upload_config.batch_size
+        ):
+            with self.connection() as conn:
+                values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
+                if self.connection_config.db_type == DatabaseType.SQLITE:
+                    conn.executemany(stmt, values)
+                else:
+                    with conn.cursor() as cur:
+                        cur.executemany(stmt, values)
+
+                conn.commit()
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        for content in contents:
+            self.upload_contents(content=content)
+
+
+sql_destination_entry = DestinationRegistryEntry(
+    connection_config=SimpleSqlConfig,
+    uploader=SQLUploader,
+    uploader_config=SQLUploaderConfig,
+    upload_stager=SQLUploadStager,
+    upload_stager_config=SQLUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/utils.py

@@ -0,0 +1,19 @@
+from datetime import datetime
+from typing import Union
+
+from dateutil import parser
+
+
+def parse_datetime(date_value: Union[int, str, float, datetime]) -> datetime:
+    if isinstance(date_value, datetime):
+        return date_value
+    elif isinstance(date_value, float):
+        return datetime.fromtimestamp(date_value)
+    elif isinstance(date_value, int):
+        return datetime.fromtimestamp(date_value / 1000)
+
+    try:
+        timestamp = float(date_value)
+        return datetime.fromtimestamp(timestamp)
+    except ValueError:
+        return parser.parse(date_value)

unstructured_ingest/v2/processes/connectors/weaviate.py

@@ -0,0 +1,235 @@
+import json
+from dataclasses import dataclass, field
+from datetime import date, datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+from dateutil import parser
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from weaviate import Client
+
+CONNECTOR_TYPE = "weaviate"
+
+
+@dataclass
+class WeaviateAccessConfig(AccessConfig):
+    access_token: Optional[str] = None
+    api_key: Optional[str] = None
+    client_secret: Optional[str] = None
+    password: Optional[str] = None
+
+
+@dataclass
+class WeaviateConnectionConfig(ConnectionConfig):
+    host_url: str
+    class_name: str
+    access_config: WeaviateAccessConfig = enhanced_field(sensitive=True)
+    username: Optional[str] = None
+    anonymous: bool = False
+    scope: Optional[list[str]] = None
+    refresh_token: Optional[str] = None
+    connector_type: str = CONNECTOR_TYPE
+
+
+@dataclass
+class WeaviateUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class WeaviateUploadStager(UploadStager):
+    upload_stager_config: WeaviateUploadStagerConfig = field(
+        default_factory=lambda: WeaviateUploadStagerConfig()
+    )
+
+    @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)
+
+    @classmethod
+    def conform_dict(cls, data: dict) -> None:
+        """
+        Updates the element dictionary to conform to the Weaviate schema
+        """
+
+        # Dict as string formatting
+        if record_locator := data.get("metadata", {}).get("data_source", {}).get("record_locator"):
+            # Explicit casting otherwise fails schema type checking
+            data["metadata"]["data_source"]["record_locator"] = str(json.dumps(record_locator))
+
+        # Array of items as string formatting
+        if points := data.get("metadata", {}).get("coordinates", {}).get("points"):
+            data["metadata"]["coordinates"]["points"] = str(json.dumps(points))
+
+        if links := data.get("metadata", {}).get("links", {}):
+            data["metadata"]["links"] = str(json.dumps(links))
+
+        if permissions_data := (
+            data.get("metadata", {}).get("data_source", {}).get("permissions_data")
+        ):
+            data["metadata"]["data_source"]["permissions_data"] = json.dumps(permissions_data)
+
+        # Datetime formatting
+        if date_created := data.get("metadata", {}).get("data_source", {}).get("date_created"):
+            data["metadata"]["data_source"]["date_created"] = cls.parse_date_string(
+                date_created
+            ).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if date_modified := data.get("metadata", {}).get("data_source", {}).get("date_modified"):
+            data["metadata"]["data_source"]["date_modified"] = cls.parse_date_string(
+                date_modified
+            ).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if date_processed := data.get("metadata", {}).get("data_source", {}).get("date_processed"):
+            data["metadata"]["data_source"]["date_processed"] = cls.parse_date_string(
+                date_processed
+            ).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if last_modified := data.get("metadata", {}).get("last_modified"):
+            data["metadata"]["last_modified"] = cls.parse_date_string(last_modified).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        # String casting
+        if version := data.get("metadata", {}).get("data_source", {}).get("version"):
+            data["metadata"]["data_source"]["version"] = str(version)
+
+        if page_number := data.get("metadata", {}).get("page_number"):
+            data["metadata"]["page_number"] = str(page_number)
+
+        if regex_metadata := data.get("metadata", {}).get("regex_metadata"):
+            data["metadata"]["regex_metadata"] = str(json.dumps(regex_metadata))
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        with open(elements_filepath) as elements_file:
+            elements_contents = json.load(elements_file)
+        for element in elements_contents:
+            self.conform_dict(data=element)
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        with open(output_path, "w") as output_file:
+            json.dump(elements_contents, output_file)
+        return output_path
+
+
+@dataclass
+class WeaviateUploaderConfig(UploaderConfig):
+    batch_size: int = 100
+
+
+@dataclass
+class WeaviateUploader(Uploader):
+    upload_config: WeaviateUploaderConfig
+    connection_config: WeaviateConnectionConfig
+    client: Optional["Client"] = field(init=False)
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["weaviate"], extras="weaviate")
+    def __post_init__(self):
+        from weaviate import Client
+
+        auth = self._resolve_auth_method()
+        self.client = Client(url=self.connection_config.host_url, auth_client_secret=auth)
+
+    @requires_dependencies(["weaviate"], extras="weaviate")
+    def _resolve_auth_method(self):
+        access_configs = self.connection_config.access_config
+        connection_config = self.connection_config
+        if connection_config.anonymous:
+            return None
+
+        if access_configs.access_token:
+            from weaviate.auth import AuthBearerToken
+
+            return AuthBearerToken(
+                access_token=access_configs.access_token,
+                refresh_token=connection_config.refresh_token,
+            )
+        elif access_configs.api_key:
+            from weaviate.auth import AuthApiKey
+
+            return AuthApiKey(api_key=access_configs.api_key)
+        elif access_configs.client_secret:
+            from weaviate.auth import AuthClientCredentials
+
+            return AuthClientCredentials(
+                client_secret=access_configs.client_secret, scope=connection_config.scope
+            )
+        elif connection_config.username and access_configs.password:
+            from weaviate.auth import AuthClientPassword
+
+            return AuthClientPassword(
+                username=connection_config.username,
+                password=access_configs.password,
+                scope=connection_config.scope,
+            )
+        return None
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        # TODO update to use async support in weaviate client
+        #  once the version can be bumped to include it
+        elements_dict = []
+        for content in contents:
+            with open(content.path) as elements_file:
+                elements = json.load(elements_file)
+                elements_dict.extend(elements)
+
+        logger.info(
+            f"writing {len(elements_dict)} objects to destination "
+            f"class {self.connection_config.class_name} "
+            f"at {self.connection_config.host_url}",
+        )
+
+        self.client.batch.configure(batch_size=self.upload_config.batch_size)
+        with self.client.batch as b:
+            for e in elements_dict:
+                vector = e.pop("embeddings", None)
+                b.add_data_object(
+                    e,
+                    self.connection_config.class_name,
+                    vector=vector,
+                )
+
+
+weaviate_destination_entry = DestinationRegistryEntry(
+    connection_config=WeaviateConnectionConfig,
+    uploader=WeaviateUploader,
+    uploader_config=WeaviateUploaderConfig,
+    upload_stager=WeaviateUploadStager,
+    upload_stager_config=WeaviateUploadStagerConfig,
+)

unstructured_ingest/v2/processes/embedder.py

@@ -0,0 +1,76 @@
+from abc import ABC
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+from unstructured.documents.elements import Element
+from unstructured.embed.interfaces import BaseEmbeddingEncoder
+from unstructured.staging.base import elements_from_json
+
+from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin, enhanced_field
+from unstructured_ingest.v2.interfaces.process import BaseProcess
+
+
+@dataclass
+class EmbedderConfig(EnhancedDataClassJsonMixin):
+    embedding_provider: Optional[str] = None
+    embedding_api_key: Optional[str] = enhanced_field(default=None, sensitive=True)
+    embedding_model_name: Optional[str] = None
+    embedding_aws_access_key_id: Optional[str] = None
+    embedding_aws_secret_access_key: Optional[str] = None
+    embedding_aws_region: Optional[str] = None
+
+    def get_embedder(self) -> BaseEmbeddingEncoder:
+        kwargs: dict[str, Any] = {}
+        if self.embedding_api_key:
+            kwargs["api_key"] = self.embedding_api_key
+        if self.embedding_model_name:
+            kwargs["model_name"] = self.embedding_model_name
+        # TODO make this more dynamic to map to encoder configs
+        if self.embedding_provider == "langchain-openai":
+            from unstructured.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+
+            return OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(**kwargs))
+        elif self.embedding_provider == "langchain-huggingface":
+            from unstructured.embed.huggingface import (
+                HuggingFaceEmbeddingConfig,
+                HuggingFaceEmbeddingEncoder,
+            )
+
+            return HuggingFaceEmbeddingEncoder(config=HuggingFaceEmbeddingConfig(**kwargs))
+        elif self.embedding_provider == "octoai":
+            from unstructured.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
+
+            return OctoAIEmbeddingEncoder(config=OctoAiEmbeddingConfig(**kwargs))
+        elif self.embedding_provider == "langchain-aws-bedrock":
+            from unstructured.embed.bedrock import BedrockEmbeddingConfig, BedrockEmbeddingEncoder
+
+            return BedrockEmbeddingEncoder(
+                config=BedrockEmbeddingConfig(
+                    aws_access_key_id=self.embedding_aws_access_key_id,
+                    aws_secret_access_key=self.embedding_aws_secret_access_key,
+                    region_name=self.embedding_aws_region,
+                )
+            )
+        elif self.embedding_provider == "langchain-vertexai":
+            from unstructured.embed.vertexai import (
+                VertexAIEmbeddingConfig,
+                VertexAIEmbeddingEncoder,
+            )
+
+            return VertexAIEmbeddingEncoder(config=VertexAIEmbeddingConfig(**kwargs))
+        else:
+            raise ValueError(f"{self.embedding_provider} not a recognized encoder")
+
+
+@dataclass
+class Embedder(BaseProcess, ABC):
+    config: EmbedderConfig
+
+    def run(self, elements_filepath: Path, **kwargs: Any) -> list[Element]:
+        # TODO update base embedder classes to support async
+        embedder = self.config.get_embedder()
+        elements = elements_from_json(filename=str(elements_filepath))
+        if not elements:
+            return elements
+        return embedder.embed_documents(elements=elements)