unstructured-ingest 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

unstructured_ingest/embed/togetherai.py

@@ -9,6 +9,7 @@ from unstructured_ingest.embed.interfaces import (
     EmbeddingConfig,
 )
 from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.errors import (
     RateLimitError as CustomRateLimitError,
@@ -71,13 +72,18 @@ class TogetherAIEmbeddingEncoder(BaseEmbeddingEncoder):
 
     def _embed_documents(self, elements: list[str]) -> list[list[float]]:
         client = self.config.get_client()
+        embeddings = []
         try:
-            outputs = client.embeddings.create(
-                model=self.config.embedder_model_name, input=elements
-            )
+            for batch in batch_generator(
+                elements, batch_size=self.config.batch_size or len(elements)
+            ):
+                outputs = client.embeddings.create(
+                    model=self.config.embedder_model_name, input=batch
+                )
+                embeddings.extend([outputs.data[i].embedding for i in range(len(batch))])
         except Exception as e:
             raise self.wrap_error(e=e)
-        return [outputs.data[i].embedding for i in range(len(elements))]
+        return embeddings
 
 
 @dataclass
@@ -97,10 +103,15 @@ class AsyncTogetherAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
 
     async def _embed_documents(self, elements: list[str]) -> list[list[float]]:
         client = self.config.get_async_client()
+        embeddings = []
         try:
-            outputs = await client.embeddings.create(
-                model=self.config.embedder_model_name, input=elements
-            )
+            for batch in batch_generator(
+                elements, batch_size=self.config.batch_size or len(elements)
+            ):
+                outputs = await client.embeddings.create(
+                    model=self.config.embedder_model_name, input=batch
+                )
+                embeddings.extend([outputs.data[i].embedding for i in range(len(batch))])
         except Exception as e:
             raise self.wrap_error(e=e)
-        return [outputs.data[i].embedding for i in range(len(elements))]
+        return embeddings

unstructured_ingest/embed/vertexai.py

@@ -13,6 +13,7 @@ from unstructured_ingest.embed.interfaces import (
     BaseEmbeddingEncoder,
     EmbeddingConfig,
 )
+from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.errors import UserAuthError
 
@@ -86,12 +87,15 @@ class VertexAIEmbeddingEncoder(BaseEmbeddingEncoder):
         from vertexai.language_models import TextEmbeddingInput
 
         inputs = [TextEmbeddingInput(text=element) for element in elements]
+        client = self.config.get_client()
+        embeddings = []
         try:
-            client = self.config.get_client()
-            embeddings = client.get_embeddings(inputs)
+            for batch in batch_generator(inputs, batch_size=self.config.batch_size or len(inputs)):
+                response = client.get_embeddings(batch)
+                embeddings.extend([e.values for e in response])
         except Exception as e:
             raise self.wrap_error(e=e)
-        return [e.values for e in embeddings]
+        return embeddings
 
 
 @dataclass
@@ -118,9 +122,12 @@ class AsyncVertexAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
         from vertexai.language_models import TextEmbeddingInput
 
         inputs = [TextEmbeddingInput(text=element) for element in elements]
+        client = self.config.get_client()
+        embeddings = []
         try:
-            client = self.config.get_client()
-            embeddings = await client.get_embeddings_async(inputs)
+            for batch in batch_generator(inputs, batch_size=self.config.batch_size or len(inputs)):
+                response = await client.get_embeddings_async(batch)
+                embeddings.extend([e.values for e in response])
         except Exception as e:
             raise self.wrap_error(e=e)
-        return [e.values for e in embeddings]
+        return embeddings

unstructured_ingest/embed/voyageai.py

@@ -9,6 +9,7 @@ from unstructured_ingest.embed.interfaces import (
     EmbeddingConfig,
 )
 from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.errors import (
     ProviderError,
@@ -25,9 +26,13 @@ if TYPE_CHECKING:
 
 
 class VoyageAIEmbeddingConfig(EmbeddingConfig):
+    batch_size: int = Field(
+        default=32,
+        le=128,
+        description="Batch size for embedding requests. VoyageAI has a limit of 128.",
+    )
     api_key: SecretStr
     embedder_model_name: str = Field(default="voyage-3", alias="model_name")
-    batch_size: Optional[int] = Field(default=None)
     truncation: Optional[bool] = Field(default=None)
     max_retries: int = 0
     timeout_in_seconds: Optional[int] = None
@@ -91,12 +96,15 @@ class VoyageAIEmbeddingEncoder(BaseEmbeddingEncoder):
         return self.config.wrap_error(e=e)
 
     def _embed_documents(self, elements: list[str]) -> list[list[float]]:
-        client: VoyageAIClient = self.config.get_client()
+        client = self.config.get_client()
+        embeddings = []
         try:
-            response = client.embed(texts=elements, model=self.config.embedder_model_name)
+            for batch in batch_generator(elements, batch_size=self.config.batch_size):
+                response = client.embed(texts=batch, model=self.config.embedder_model_name)
+                embeddings.extend(response.embeddings)
         except Exception as e:
             raise self.wrap_error(e=e)
-        return response.embeddings
+        return embeddings
 
     def embed_documents(self, elements: list[dict]) -> list[dict]:
         embeddings = self._embed_documents([e.get("text", "") for e in elements])
@@ -115,11 +123,16 @@ class AsyncVoyageAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
 
     async def _embed_documents(self, elements: list[str]) -> list[list[float]]:
         client = self.config.get_async_client()
+        embeddings = []
         try:
-            response = await client.embed(texts=elements, model=self.config.embedder_model_name)
+            for batch in batch_generator(
+                elements, batch_size=self.config.batch_size or len(elements)
+            ):
+                response = await client.embed(texts=batch, model=self.config.embedder_model_name)
+                embeddings.extend(response.embeddings)
         except Exception as e:
             raise self.wrap_error(e=e)
-        return response.embeddings
+        return embeddings
 
     async def embed_documents(self, elements: list[dict]) -> list[dict]:
         embeddings = await self._embed_documents([e.get("text", "") for e in elements])

unstructured_ingest/utils/html.py

@@ -1,109 +1,159 @@
 import base64
 from pathlib import Path
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 from urllib.parse import urlparse
 from uuid import NAMESPACE_DNS, uuid5
 
-import requests
-from bs4 import BeautifulSoup
-from requests import Session
+from pydantic import BaseModel, Field
 
+from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import DownloadResponse, FileData, SourceIdentifiers
 from unstructured_ingest.v2.logger import logger
 
+if TYPE_CHECKING:
+    from bs4.element import Tag
+    from requests import Session
 
-def convert_image_tags(url: str, original_html: str, session: Optional[Session] = None) -> str:
-    session = session or requests.Session()
-    parsed_url = urlparse(url)
-    base_url = parsed_url.scheme + "://" + parsed_url.netloc
-    soup = BeautifulSoup(original_html, "html.parser")
-    images = soup.find_all("img")
-    for image in images:
-        current_source = image["src"]
-        if current_source.startswith("//"):
-            source_url = f"{parsed_url.scheme}:{current_source}"
-        elif current_source.startswith("http"):
-            source_url = current_source
+
+class HtmlMixin(BaseModel):
+    extract_images: bool = Field(
+        default=False,
+        description="if true, will download images and replace "
+        "the html content with base64 encoded images",
+    )
+    extract_files: bool = Field(
+        default=False, description="if true, will download any embedded files"
+    )
+    force_download: bool = Field(
+        default=False,
+        description="if true, will redownload extracted files even if they already exist locally",
+    )
+    allow_list: Optional[list[str]] = Field(
+        default=None,
+        description="list of allowed urls to download, if not set, "
+        "will default to the base url the original HTML came from",
+    )
+
+    @requires_dependencies(["requests"])
+    def get_default_session(self) -> "Session":
+        import requests
+
+        return requests.Session()
+
+    def get_absolute_url(self, tag_link: str, url: str) -> str:
+        parsed_url = urlparse(url)
+        base_url = parsed_url.scheme + "://" + parsed_url.netloc
+        if tag_link.startswith("//"):
+            return f"{parsed_url.scheme}:{tag_link}"
+        elif tag_link.startswith("http"):
+            return tag_link
         else:
-            source_url = base_url + current_source
-        try:
-            response = session.get(source_url)
-            response.raise_for_status()
-            image_content = response.content
-            logger.debug(
-                "img tag having src updated from {} to base64 content".format(image["src"])
-            )
-            image["src"] = f"data:image/png;base64,{base64.b64encode(image_content).decode()}"
-        except Exception as e:
-            logger.warning(
-                f"failed to download image content from {source_url}: {e}", exc_info=True
-            )
-    return str(soup)
-
-
-def download_link(
-    download_dir: Path, link: str, session: Optional[Session] = None, force_download: bool = False
-) -> Path:
-    session = session or requests.Session()
-    filename = Path(urlparse(url=link).path).name
-    download_path = download_dir / filename
-    logger.debug(f"downloading file from {link} to {download_path}")
-    if download_path.exists() and download_path.is_file() and not force_download:
-        return download_path
-    with download_path.open("wb") as downloaded_file:
-        response = session.get(link)
+            tag_link = tag_link.lstrip("/")
+            return f"{base_url}/{tag_link}"
+
+    def download_content(self, url: str, session: "Session") -> bytes:
+        response = session.get(url)
         response.raise_for_status()
-        downloaded_file.write(response.content)
-    return download_path
-
-
-def download_embedded_files(
-    download_dir: Path,
-    original_filedata: FileData,
-    original_html: str,
-    session: Optional[Session] = None,
-    force_download: bool = False,
-) -> list[DownloadResponse]:
-    session = session or requests.Session()
-    url = original_filedata.metadata.url
-    parsed_url = urlparse(url)
-    base_url = parsed_url.scheme + "://" + parsed_url.netloc
-    soup = BeautifulSoup(original_html, "html.parser")
-    tags = soup.find_all("a", href=True)
-    hrefs = [
-        tag["href"]
-        for tag in tags
-        if not tag["href"].startswith("#") and Path(tag["href"]).suffix != ""
-    ]
-    results = []
-    for current_source in hrefs:
-        download_dir.mkdir(parents=True, exist_ok=True)
-        if current_source.startswith("//"):
-            source_url = f"{parsed_url.scheme}:{current_source}"
-        elif current_source.startswith("http"):
-            source_url = current_source
-        else:
-            source_url = base_url + current_source
-        try:
-            downloaded_path = download_link(
-                download_dir=download_dir,
-                link=source_url,
-                session=session,
-                force_download=force_download,
-            )
-        except Exception as e:
-            logger.warning(f"failed to download file content from {source_url}: {e}")
-            continue
-        result_file_data = original_filedata.model_copy(deep=True)
-        result_file_data.metadata.url = source_url
-        result_file_data.metadata.record_locator["parent_url"] = url
-        result_file_data.identifier = str(
-            uuid5(NAMESPACE_DNS, source_url + original_filedata.identifier)
+        return response.content
+
+    def can_download(self, url_to_download: str, original_url: str) -> bool:
+        parsed_original_url = urlparse(original_url)
+        base_url = parsed_original_url.scheme + "://" + parsed_original_url.netloc
+        allow_list = self.allow_list or [base_url]
+        for allowed_url in allow_list:
+            if url_to_download.startswith(allowed_url):
+                return True
+        logger.info(f"Skipping url because it does not match the allow list: {url_to_download}")
+        return False
+
+    def extract_image_src(self, image: "Tag", url: str, session: "Session") -> "Tag":
+        current_src = image["src"]
+        if current_src.startswith("data:image/png;base64"):
+            # already base64 encoded
+            return image
+        absolute_url = self.get_absolute_url(tag_link=image["src"], url=url)
+        if not self.can_download(url_to_download=absolute_url, original_url=url):
+            return image
+        image_content = self.download_content(url=absolute_url, session=session)
+        logger.debug("img tag having src updated from {} to base64 content".format(image["src"]))
+        image["src"] = f"data:image/png;base64,{base64.b64encode(image_content).decode()}"
+        return image
+
+    @requires_dependencies(["bs4"])
+    def extract_html_images(self, url: str, html: str, session: Optional["Session"] = None) -> str:
+        from bs4 import BeautifulSoup
+
+        session = session or self.get_default_session()
+        soup = BeautifulSoup(html, "html.parser")
+        images = soup.find_all("img")
+        for image in images:
+            self.extract_image_src(image=image, url=url, session=session)
+        return str(soup)
+
+    @requires_dependencies(["bs4"])
+    def get_hrefs(self, url: str, html: str) -> list:
+        from bs4 import BeautifulSoup
+
+        soup = BeautifulSoup(html, "html.parser")
+        tags = soup.find_all("a", href=True)
+        hrefs = [
+            tag["href"]
+            for tag in tags
+            if not tag["href"].startswith("#") and Path(tag["href"]).suffix != ""
+        ]
+        absolute_urls = [self.get_absolute_url(tag_link=href, url=url) for href in hrefs]
+        allowed_urls = [
+            url_to_download
+            for url_to_download in absolute_urls
+            if self.can_download(url_to_download=url_to_download, original_url=url)
+        ]
+        return allowed_urls
+
+    def write_content(self, content: bytes, path: Path) -> None:
+        if path.exists() and path.is_file() and not self.force_download:
+            return
+        if not path.parent.exists():
+            path.parent.mkdir(parents=True)
+        with path.open("wb") as f:
+            f.write(content)
+
+    def get_download_response(
+        self, url: str, download_dir: Path, file_data: FileData, session: "Session"
+    ) -> DownloadResponse:
+        filename = Path(urlparse(url=url).path).name
+        download_path = download_dir / filename
+        self.write_content(
+            content=self.download_content(url=url, session=session), path=download_path
         )
-        filename = Path(urlparse(url=source_url).path).name
+        result_file_data = file_data.model_copy(deep=True)
+        result_file_data.metadata.url = url
+        if result_file_data.metadata.record_locator is None:
+            result_file_data.metadata.record_locator = {}
+        result_file_data.metadata.record_locator["parent_url"] = url
+        result_file_data.identifier = str(uuid5(NAMESPACE_DNS, url + file_data.identifier))
+        filename = Path(urlparse(url=url).path).name
         result_file_data.source_identifiers = SourceIdentifiers(
             filename=filename, fullpath=filename
         )
-        result_file_data.local_download_path = downloaded_path.as_posix()
-        results.append(DownloadResponse(file_data=result_file_data, path=downloaded_path))
-    return results
+        result_file_data.local_download_path = download_path.as_posix()
+        return DownloadResponse(file_data=result_file_data, path=download_path)
+
+    def extract_embedded_files(
+        self,
+        url: str,
+        html: str,
+        download_dir: Path,
+        original_filedata: FileData,
+        session: Optional["Session"] = None,
+    ) -> list[DownloadResponse]:
+        session = session or self.get_default_session()
+        urls_to_download = self.get_hrefs(url=url, html=html)
+        return [
+            self.get_download_response(
+                url=url_to_download,
+                download_dir=download_dir,
+                file_data=original_filedata,
+                session=session,
+            )
+            for url_to_download in urls_to_download
+        ]

unstructured_ingest/v2/interfaces/__init__.py

@@ -5,7 +5,7 @@ from .indexer import Indexer, IndexerConfig
 from .process import BaseProcess
 from .processor import ProcessorConfig
 from .upload_stager import UploadStager, UploadStagerConfig
-from .uploader import UploadContent, Uploader, UploaderConfig
+from .uploader import UploadContent, Uploader, UploaderConfig, VectorDBUploader
 
 __all__ = [
     "DownloadResponse",
@@ -29,4 +29,5 @@ __all__ = [
     "FileDataSourceMetadata",
     "BatchFileData",
     "BatchItem",
+    "VectorDBUploader",
 ]

unstructured_ingest/v2/interfaces/indexer.py

@@ -1,4 +1,4 @@
-from abc import ABC, abstractmethod
+from abc import ABC
 from typing import Any, AsyncGenerator, Generator, Optional, TypeVar
 
 from pydantic import BaseModel
@@ -22,9 +22,8 @@ class Indexer(BaseProcess, BaseConnector, ABC):
     def is_async(self) -> bool:
         return False
 
-    @abstractmethod
     def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        pass
+        raise NotImplementedError()
 
     async def run_async(self, **kwargs: Any) -> AsyncGenerator[FileData, None]:
         raise NotImplementedError()

unstructured_ingest/v2/interfaces/process.py

@@ -8,6 +8,9 @@ class BaseProcess(ABC):
     def is_async(self) -> bool:
         return False
 
+    def init(self, *kwargs: Any) -> None:
+        pass
+
     def precheck(self) -> None:
         pass
 

unstructured_ingest/v2/interfaces/uploader.py

@@ -1,7 +1,7 @@
 from abc import ABC
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import Any, Optional, TypeVar
 
 from pydantic import BaseModel
 
@@ -38,6 +38,11 @@ class Uploader(BaseProcess, BaseConnector, ABC):
     def run_batch(self, contents: list[UploadContent], **kwargs: Any) -> None:
         raise NotImplementedError()
 
+    def create_destination(self, destination_name: str = "elements", **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_data(path=path)
         self.run_data(data=data, file_data=file_data, **kwargs)
@@ -51,3 +56,11 @@ class Uploader(BaseProcess, BaseConnector, ABC):
 
     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, destination_name: str = "elements", vector_length: Optional[int] = None, **kwargs: Any
+    ) -> bool:
+        return False

unstructured_ingest/v2/pipeline/pipeline.py

@@ -11,6 +11,7 @@ from typing import Any
 from unstructured_ingest.v2.interfaces import ProcessorConfig, Uploader
 from unstructured_ingest.v2.logger import logger, make_default_logger
 from unstructured_ingest.v2.otel import OtelHandler
+from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.pipeline.steps.chunk import Chunker, ChunkStep
 from unstructured_ingest.v2.pipeline.steps.download import DownloaderT, DownloadStep
 from unstructured_ingest.v2.pipeline.steps.embed import Embedder, EmbedStep
@@ -91,10 +92,6 @@ class Pipeline:
         self.chunker_step = ChunkStep(process=chunker, context=self.context) if chunker else None
 
         self.embedder_step = EmbedStep(process=embedder, context=self.context) if embedder else None
-        # TODO: support initialize() call from each step process
-        # Potential long call to download embedder models, run before any fanout:
-        if embedder and embedder.config:
-            embedder.config.get_embedder().initialize()
 
         self.stager_step = UploadStageStep(process=stager, context=self.context) if stager else None
         self.uploader_step = UploadStep(process=uploader, context=self.context)
@@ -135,6 +132,7 @@ class Pipeline:
             with otel_handler.get_tracer().start_as_current_span(
                 "ingest process", record_exception=True
             ):
+                self._run_inits()
                 self._run_prechecks()
                 self._run()
         finally:
@@ -156,7 +154,7 @@ class Pipeline:
         final = [f for f in flat if f]
         return final or None
 
-    def _run_prechecks(self):
+    def _get_all_steps(self) -> list[PipelineStep]:
         steps = [self.indexer_step, self.downloader_step, self.partitioner_step, self.uploader_step]
         if self.chunker_step:
             steps.append(self.chunker_step)
@@ -166,8 +164,24 @@ class Pipeline:
             steps.append(self.uncompress_step)
         if self.stager_step:
             steps.append(self.stager_step)
+        return steps
+
+    def _run_inits(self):
+        failures = {}
+
+        for step in self._get_all_steps():
+            try:
+                step.process.init()
+            except Exception as e:
+                failures[step.process.__class__.__name__] = f"[{type(e).__name__}] {e}"
+        if failures:
+            for k, v in failures.items():
+                logger.error(f"Step init failure: {k}: {v}")
+            raise PipelineError("Init failed")
+
+    def _run_prechecks(self):
         failures = {}
-        for step in steps:
+        for step in self._get_all_steps():
             try:
                 step.process.precheck()
             except Exception as e:

unstructured_ingest/v2/processes/connectors/confluence.py

@@ -7,6 +7,7 @@ from pydantic import Field, Secret
 
 from unstructured_ingest.error import SourceConnectionError
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.utils.html import HtmlMixin
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
@@ -180,19 +181,8 @@ class ConfluenceIndexer(Indexer):
                 yield file_data
 
 
-class ConfluenceDownloaderConfig(DownloaderConfig):
-    extract_images: bool = Field(
-        default=False,
-        description="if true, will download images and replace "
-        "the html content with base64 encoded images",
-    )
-    extract_files: bool = Field(
-        default=False, description="if true, will download any embedded files"
-    )
-    force_download: bool = Field(
-        default=False,
-        description="if true, will redownload extracted files even if they already exist locally",
-    )
+class ConfluenceDownloaderConfig(DownloaderConfig, HtmlMixin):
+    pass
 
 
 @dataclass
@@ -206,24 +196,27 @@ class ConfluenceDownloader(Downloader):
     ) -> list[DownloadResponse]:
         if not self.download_config.extract_files:
             return []
-        from unstructured_ingest.utils.html import download_embedded_files
-
+        url = current_file_data.metadata.url
+        if url is None:
+            logger.warning(
+                f"""Missing URL for file: {current_file_data.source_identifiers.filename}.
+                Skipping file extraction."""
+            )
+            return []
         filepath = current_file_data.source_identifiers.relative_path
         download_path = Path(self.download_dir) / filepath
         download_dir = download_path.with_suffix("")
-        return download_embedded_files(
+        return self.download_config.extract_embedded_files(
+            url=url,
             download_dir=download_dir,
             original_filedata=current_file_data,
-            original_html=html,
+            html=html,
             session=session,
-            force_download=self.download_config.force_download,
         )
 
     def run(self, file_data: FileData, **kwargs) -> download_responses:
         from bs4 import BeautifulSoup
 
-        from unstructured_ingest.utils.html import convert_image_tags
-
         doc_id = file_data.identifier
         try:
             with self.connection_config.get_client() as client:
@@ -246,8 +239,8 @@ class ConfluenceDownloader(Downloader):
         content = f"<body class='Document' >{title_html}{content}</body>"
         if self.download_config.extract_images:
             with self.connection_config.get_client() as client:
-                content = convert_image_tags(
-                    url=file_data.metadata.url, original_html=content, session=client._session
+                content = self.download_config.extract_html_images(
+                    url=file_data.metadata.url, html=content, session=client._session
                 )
 
         filepath = file_data.source_identifiers.relative_path