unstructured-ingest 1.2.32__py3-none-any.whl

unstructured_ingest/interfaces/upload_stager.py

@@ -0,0 +1,89 @@
+from abc import ABC
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+
+from unstructured_ingest.data_types.file_data import FileData
+from unstructured_ingest.interfaces import BaseProcess
+from unstructured_ingest.utils import ndjson
+from unstructured_ingest.utils.data_prep import get_json_data, write_data
+
+
+class UploadStagerConfig(BaseModel):
+    pass
+
+
+UploadStagerConfigT = TypeVar("UploadStagerConfigT", bound=UploadStagerConfig)
+
+
+@dataclass
+class UploadStager(BaseProcess, ABC):
+    upload_stager_config: UploadStagerConfigT
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        return element_dict
+
+    def get_output_path(self, output_filename: str, output_dir: Path) -> Path:
+        output_path = Path(output_filename)
+        output_filename = f"{Path(output_filename).stem}{output_path.suffix}"
+        output_path = Path(output_dir) / Path(f"{output_filename}")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        return output_path
+
+    def stream_update(self, input_file: Path, output_file: Path, file_data: FileData) -> None:
+        with input_file.open() as in_f:
+            reader = ndjson.reader(in_f)
+            with output_file.open("w") as out_f:
+                writer = ndjson.writer(out_f)
+                for element in reader:
+                    conformed_element = self.conform_dict(element_dict=element, file_data=file_data)
+                    writer.write(row=conformed_element)
+                    writer.f.flush()
+
+    def process_whole(self, input_file: Path, output_file: Path, file_data: FileData) -> None:
+        elements_contents = get_json_data(path=input_file)
+
+        conformed_elements = [
+            self.conform_dict(element_dict=element, file_data=file_data)
+            for element in elements_contents
+        ]
+        write_data(path=output_file, data=conformed_elements)
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        output_file = self.get_output_path(output_filename=output_filename, output_dir=output_dir)
+        if elements_filepath.suffix == ".ndjson":
+            self.stream_update(
+                input_file=elements_filepath, output_file=output_file, file_data=file_data
+            )
+        elif elements_filepath.suffix == ".json":
+            self.process_whole(
+                input_file=elements_filepath, output_file=output_file, file_data=file_data
+            )
+        else:
+            raise ValueError(f"Unsupported file extension: {elements_filepath}")
+        return output_file
+
+    async def run_async(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        return self.run(
+            elements_filepath=elements_filepath,
+            output_dir=output_dir,
+            output_filename=output_filename,
+            file_data=file_data,
+            **kwargs,
+        )

unstructured_ingest/interfaces/uploader.py

@@ -0,0 +1,67 @@
+from abc import ABC
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+
+from unstructured_ingest.data_types.file_data import FileData
+from unstructured_ingest.interfaces import BaseConnector, BaseProcess
+from unstructured_ingest.utils.data_prep import get_json_data
+
+
+class UploaderConfig(BaseModel):
+    pass
+
+
+UploaderConfigT = TypeVar("UploaderConfigT", bound=UploaderConfig)
+
+
+@dataclass
+class UploadContent:
+    path: Path
+    file_data: FileData
+
+
+@dataclass
+class Uploader(BaseProcess, BaseConnector, ABC):
+    upload_config: UploaderConfigT
+    connector_type: str
+
+    def is_async(self) -> bool:
+        return False
+
+    def is_batch(self) -> bool:
+        return False
+
+    def run_batch(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        raise NotImplementedError()
+
+    def create_destination(
+        self, destination_name: str = "unstructuredautocreated", **kwargs: Any
+    ) -> bool:
+        # Update the uploader config if needed with a new destination that gets created.
+        # Return a flag on if anything was created or not.
+        return False
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        data = get_json_data(path=path)
+        self.run_data(data=data, file_data=file_data, **kwargs)
+
+    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        data = get_json_data(path=path)
+        await self.run_data_async(data=data, file_data=file_data, **kwargs)
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        raise NotImplementedError()
+
+    async def run_data_async(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        return self.run_data(data=data, file_data=file_data, **kwargs)
+
+
+@dataclass
+class VectorDBUploader(Uploader, ABC):
+    def create_destination(
+        self, vector_length: int, destination_name: str = "unstructuredautocreated", **kwargs: Any
+    ) -> bool:
+        return False

unstructured_ingest/logger.py

@@ -0,0 +1,39 @@
+import logging
+
+logger = logging.getLogger("unstructured_ingest")
+
+
+def remove_root_handlers(logger: logging.Logger) -> None:
+    # NOTE(robinson): in some environments such as Google Colab, there is a root handler
+    # that doesn't not mask secrets, meaning sensitive info such as api keys appear in logs.
+    # Removing these when they exist prevents this behavior
+    if logger.root.hasHandlers():
+        for handler in logger.root.handlers:
+            logger.root.removeHandler(handler)
+
+
+def ingest_log_streaming_init(level: int) -> None:
+    handler = logging.StreamHandler()
+    handler.name = "ingest_log_handler"
+    formatter = logging.Formatter("%(asctime)s %(processName)-10s %(levelname)-8s %(message)s")
+    handler.setFormatter(formatter)
+
+    # Only want to add the handler once
+    if "ingest_log_handler" not in [h.name for h in logger.handlers]:
+        logger.addHandler(handler)
+
+    remove_root_handlers(logger)
+    logger.setLevel(level)
+
+
+def make_default_logger(level: int) -> logging.Logger:
+    """Return a custom logger."""
+    logger = logging.getLogger("unstructured_ingest")
+    handler = logging.StreamHandler()
+    handler.name = "ingest_log_handler"
+    formatter = logging.Formatter("%(asctime)s %(processName)-10s %(levelname)-8s %(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel(level)
+    remove_root_handlers(logger)
+    return logger

unstructured_ingest/main.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+from unstructured_ingest.cli.cli import get_cmd
+
+
+def main():
+    ingest_cmd = get_cmd()
+    ingest_cmd()
+
+
+if __name__ == "__main__":
+    main()

unstructured_ingest/otel.py

@@ -0,0 +1,128 @@
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import Callable, ClassVar, Optional, Protocol, Sequence
+
+from opentelemetry import trace
+from opentelemetry.context import attach, get_current
+from opentelemetry.propagate import extract, inject
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import ReadableSpan, Tracer, TracerProvider
+from opentelemetry.sdk.trace.export import (
+    ConsoleSpanExporter,
+    SimpleSpanProcessor,
+    SpanExportResult,
+)
+
+from unstructured_ingest.logger import logger
+
+
+class AddTraceCallable(Protocol):
+    def __call__(self, provider: TracerProvider) -> None:
+        pass
+
+
+class LogSpanExporter(ConsoleSpanExporter):
+    def __init__(self, log_out: Callable = logger.info, **kwargs):
+        self.log_out = log_out
+        super().__init__(**kwargs)
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        for span in spans:
+            self.log_out(self.formatter(span))
+        return SpanExportResult.SUCCESS
+
+
+def get_log_out() -> Callable:
+    level_names_mapping = {
+        "CRITICAL": logging.CRITICAL,
+        "FATAL": logging.FATAL,
+        "ERROR": logging.ERROR,
+        "WARN": logging.WARNING,
+        "WARNING": logging.WARNING,
+        "INFO": logging.INFO,
+        "DEBUG": logging.DEBUG,
+        "NOTSET": logging.NOTSET,
+    }
+    log_level = os.getenv("OTEL_LOG_LEVEL", "DEBUG").upper()
+    log_level_int = level_names_mapping.get(log_level, logging.DEBUG)
+    return lambda message: logger.log(log_level_int, message)
+
+
+@dataclass
+class OtelHandler:
+    otel_endpoint: Optional[str] = None
+    service_name: str = "unstructured-ingest"
+    trace_provider: TracerProvider = field(init=False)
+    log_out: Callable = field(default=get_log_out())
+    trace_context_key: ClassVar[str] = "_trace_context"
+
+    def init_trace(self):
+        # Should only be done once
+        resource = Resource(attributes={SERVICE_NAME: self.service_name})
+        trace_provider = self.init_trace_provider(resource=resource)
+        trace.set_tracer_provider(trace_provider)
+
+    @staticmethod
+    def set_attributes(span, attributes_dict):
+        if attributes_dict:
+            for att in attributes_dict:
+                span.set_attribute(att, attributes_dict[att])
+
+    @staticmethod
+    def inject_context() -> dict:
+        trace_context = {}
+        current_context = get_current()
+        inject(trace_context, current_context)
+        return trace_context
+
+    @staticmethod
+    def attach_context(trace_context: dict) -> object:
+        extracted_context = extract(trace_context)
+        return attach(extracted_context)
+
+    def get_otel_endpoint(self) -> Optional[str]:
+        if otel_endpoint := self.otel_endpoint:
+            return otel_endpoint
+        if otlp_endpoint := os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT"):
+            return otlp_endpoint
+        if otlp_traces_endpoint := os.getenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"):
+            return otlp_traces_endpoint
+        return None
+
+    def _add_console_trace_processor(self, provider: TracerProvider) -> None:
+        def custom_formatter(span: ReadableSpan) -> str:
+            duration = (span.end_time - span.start_time) / 1e9
+            s = f"{span.name} finished in {duration}s"
+            if span.attributes:
+                attributes_str = ", ".join([f"{k}={v}" for k, v in span.attributes.items()])
+                s += f", attributes: {attributes_str}"
+            return s
+
+        tracer_exporter = LogSpanExporter(formatter=custom_formatter, log_out=self.log_out)
+        processor = SimpleSpanProcessor(tracer_exporter)
+        provider.add_span_processor(span_processor=processor)
+
+    def _add_otel_trace_processor(self, provider: TracerProvider) -> None:
+        otel_endpoint = self.get_otel_endpoint()
+        if not otel_endpoint:
+            return None
+        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+
+        logger.debug(f"adding otel exported at {otel_endpoint}")
+        trace_exporter = OTLPSpanExporter()
+        processor = SimpleSpanProcessor(trace_exporter)
+        provider.add_span_processor(processor)
+
+    def init_trace_provider(self, resource: Resource) -> TracerProvider:
+        trace_provider = TracerProvider(resource=resource)
+        add_fns: list[AddTraceCallable] = [
+            self._add_otel_trace_processor,
+            self._add_console_trace_processor,
+        ]
+        for add_fn in add_fns:
+            add_fn(provider=trace_provider)
+        return trace_provider
+
+    def get_tracer(self) -> Tracer:
+        return trace.get_tracer(self.service_name)

unstructured_ingest/pipeline/interfaces.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import multiprocessing as mp
+import shutil
+from abc import ABC, abstractmethod
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Optional, TypeVar
+
+from tqdm import tqdm
+from tqdm.asyncio import tqdm as tqdm_asyncio
+
+from unstructured_ingest.interfaces import BaseProcess, ProcessorConfig, Uploader
+from unstructured_ingest.logger import logger, make_default_logger
+from unstructured_ingest.otel import OtelHandler
+from unstructured_ingest.pipeline.otel import instrument
+
+BaseProcessT = TypeVar("BaseProcessT", bound=BaseProcess)
+iterable_input = list[dict[str, Any]]
+
+
+@dataclass
+class PipelineStep(ABC):
+    process: BaseProcessT
+    context: ProcessorConfig
+    identifier: str
+
+    def __str__(self):
+        return self.identifier
+
+    def process_serially(self, iterable: iterable_input) -> Any:
+        logger.info("processing content serially")
+        if iterable:
+            if len(iterable) == 1:
+                return [self.run(**iterable[0])]
+            if self.context.tqdm:
+                return [self.run(**it) for it in tqdm(iterable, desc=self.identifier)]
+            return [self.run(**it) for it in iterable]
+        return [self.run()]
+
+    async def _process_async(self, iterable: iterable_input) -> Any:
+        if iterable:
+            if len(iterable) == 1:
+                return [await self.run_async(**iterable[0])]
+            if self.context.tqdm:
+                return await tqdm_asyncio.gather(
+                    *[self.run_async(**i) for i in iterable], desc=self.identifier
+                )
+            return await asyncio.gather(*[self.run_async(**i) for i in iterable])
+        return [await self.run_async()]
+
+    def process_async(self, iterable: iterable_input) -> Any:
+        logger.info("processing content async")
+        return self.asyncio_run(fn=self._process_async, iterable=iterable)
+
+    def asyncio_run(
+        self, fn: Callable[[Any, Any], Awaitable[Any]], *args: Any, **kwargs: Any
+    ) -> Any:
+        current_loop = asyncio._get_running_loop()
+        if current_loop is None:
+            return asyncio.run(fn(*args, **kwargs))
+        with ThreadPoolExecutor(thread_name_prefix="asyncio") as thread_pool:
+            logger.warning(
+                f"async code being run in dedicated thread pool "
+                f"to not conflict with existing event loop: {current_loop}"
+            )
+
+            def wrapped():
+                return asyncio.run(fn(*args, **kwargs))
+
+            future = thread_pool.submit(wrapped)
+            return future.result()
+
+    def process_multiprocess(self, iterable: iterable_input) -> Any:
+        logger.info("processing content across processes")
+
+        if iterable:
+            if len(iterable) == 1:
+                return self.process_serially(iterable)
+            if self.context.num_processes == 1:
+                return self.process_serially(iterable)
+            with mp.Pool(
+                processes=self.context.num_processes,
+                initializer=self._init_mp,
+                initargs=(
+                    logging.DEBUG if self.context.verbose else logging.INFO,
+                    self.context.otel_endpoint,
+                ),
+            ) as pool:
+                otel_context = OtelHandler.inject_context()
+                for iter in iterable:
+                    iter[OtelHandler.trace_context_key] = otel_context
+                if self.context.tqdm:
+                    return list(
+                        tqdm(
+                            pool.imap_unordered(func=self._wrap_mp, iterable=iterable),
+                            total=len(iterable),
+                            desc=self.identifier,
+                        )
+                    )
+                return pool.map(self._wrap_mp, iterable)
+        return [self.run()]
+
+    def _wrap_mp(self, input_kwargs: dict) -> Any:
+        # Allow mapping of kwargs via multiprocessing map()
+        return self.run(**input_kwargs)
+
+    def _init_mp(self, log_level: int, endpoint: Optional[str] = None) -> None:
+        # Init logger for each spawned process when using multiprocessing pool
+        make_default_logger(level=log_level)
+        otel_handler = OtelHandler(otel_endpoint=endpoint, log_out=logger.debug)
+        otel_handler.init_trace()
+
+    @instrument()
+    def __call__(self, iterable: Optional[iterable_input] = None) -> Any:
+        iterable = iterable or []
+        if iterable:
+            logger.info(
+                f"calling {self.__class__.__name__} with {len(iterable)} docs",  # type: ignore
+            )
+        else:
+            logger.info(f"calling {self.__class__.__name__} with no inputs")
+        if self.context.async_supported and self.process.is_async():
+            return self.process_async(iterable=iterable)
+        if self.context.mp_supported:
+            return self.process_multiprocess(iterable=iterable)
+        return self.process_serially(iterable=iterable)
+
+    def _run(self, fn: Callable, **kwargs: Any) -> Optional[Any]:
+        return self.asyncio_run(fn=self.run_async, _fn=fn, **kwargs)
+
+    async def _run_async(self, fn: Callable, **kwargs: Any) -> Optional[Any]:
+        raise NotImplementedError
+
+    def run(self, _fn: Callable[..., Any] | None = None, **kwargs: Any) -> Optional[Any]:
+        kwargs = kwargs.copy()
+        otel_handler = OtelHandler(otel_endpoint=self.context.otel_endpoint, log_out=logger.debug)
+        tracer = otel_handler.get_tracer()
+        if trace_context := kwargs.pop(otel_handler.trace_context_key, {}):
+            otel_handler.attach_context(trace_context=trace_context)
+        attributes = {}
+        if file_data_path := kwargs.get("file_data_path"):
+            attributes["file_id"] = Path(file_data_path).stem
+        try:
+            with tracer.start_as_current_span(self.identifier, record_exception=True) as span:
+                otel_handler.set_attributes(span, attributes)
+                fn = _fn or self.process.run
+                return self._run(fn=fn, **kwargs)
+        except Exception as e:
+            logger.error(f"Exception raised while running {self.identifier}", exc_info=e)
+            if "file_data_path" in kwargs:
+                self.context.status[kwargs["file_data_path"]] = {self.identifier: str(e)}
+            if self.context.raise_on_error:
+                raise e
+            return None
+
+    async def run_async(self, _fn: Optional[Callable] = None, **kwargs: Any) -> Optional[Any]:
+        otel_handler = OtelHandler(otel_endpoint=self.context.otel_endpoint, log_out=logger.debug)
+        try:
+            attributes = {}
+            if file_data_path := kwargs.get("file_data_path"):
+                attributes["file_id"] = Path(file_data_path).stem
+            with otel_handler.get_tracer().start_as_current_span(
+                self.identifier, record_exception=True
+            ) as span:
+                otel_handler.set_attributes(span, attributes)
+                fn = _fn or self.process.run_async
+                return await self._run_async(fn=fn, **kwargs)
+        except Exception as e:
+            logger.error(f"Exception raised while running {self.identifier}", exc_info=e)
+            if "file_data_path" in kwargs:
+                self.context.status[kwargs["file_data_path"]] = {self.identifier: str(e)}
+            if self.context.raise_on_error:
+                raise e
+            return None
+
+    @property
+    def cache_dir(self) -> Path:
+        return Path(self.context.work_dir) / self.identifier
+
+    def delete_cache(self):
+        if self.context.iter_delete and self.cache_dir.exists():
+            cache_dir = self.cache_dir
+            logger.info(f"deleting {self.identifier} cache dir {cache_dir}")
+            shutil.rmtree(cache_dir)
+
+
+@dataclass
+class BatchPipelineStep(PipelineStep, ABC):
+    process: Uploader
+
+    def __call__(self, iterable: Optional[iterable_input] = None) -> Any:
+        if self.context.mp_supported and self.process.is_batch():
+            return self.run_batch(contents=iterable)
+        super().__call__(iterable=iterable)
+
+    @abstractmethod
+    def _run_batch(self, contents: iterable_input, **kwargs) -> Any:
+        pass
+
+    def run_batch(self, contents: iterable_input, **kwargs) -> Any:
+        try:
+            return self._run_batch(contents=contents, **kwargs)
+        except Exception as e:
+            self.context.status[self.identifier] = {"step_error": str(e)}
+            if self.context.raise_on_error:
+                raise e
+            return None

unstructured_ingest/pipeline/otel.py

@@ -0,0 +1,32 @@
+from functools import wraps
+from typing import Callable, Optional
+
+from unstructured_ingest.logger import logger
+from unstructured_ingest.otel import OtelHandler
+
+
+def instrument(
+    span_name: Optional[str] = None,
+    record_exception: bool = True,
+    attributes: dict[str, str] = None,
+    log_out: Callable = logger.info,
+) -> Callable[[Callable], Callable]:
+    def span_decorator(func: Callable) -> Callable:
+        def get_name(self) -> str:
+            if span_name:
+                return span_name
+            return f"{self.identifier} step"
+
+        @wraps(func)
+        def wrap_with_span(self, *args, **kwargs):
+            name = get_name(self=self)
+            otel_handler = OtelHandler(otel_endpoint=self.context.otel_endpoint, log_out=log_out)
+            with otel_handler.get_tracer().start_as_current_span(
+                name, record_exception=record_exception
+            ) as span:
+                otel_handler.set_attributes(span, attributes)
+                return func(self, *args, **kwargs)
+
+        return wrap_with_span
+
+    return span_decorator