airbyte-cdk 0.57.2__tar.gz → 0.57.4__tar.gz

{airbyte-cdk-0.57.2 → airbyte-cdk-0.57.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 0.57.2
+Version: 0.57.4
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://github.com/airbytehq/airbyte
 Author: Airbyte

{airbyte-cdk-0.57.2 → airbyte-cdk-0.57.4}/airbyte_cdk/destinations/vector_db_based/embedder.py

@@ -4,6 +4,7 @@
 
 import os
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import List, Optional, Union, cast
 
 from airbyte_cdk.destinations.vector_db_based.config import (
@@ -15,8 +16,8 @@ from airbyte_cdk.destinations.vector_db_based.config import (
     OpenAIEmbeddingConfigModel,
     ProcessingConfigModel,
 )
-from airbyte_cdk.destinations.vector_db_based.document_processor import Chunk
 from airbyte_cdk.destinations.vector_db_based.utils import create_chunks, format_exception
+from airbyte_cdk.models import AirbyteRecordMessage
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException, FailureType
 from langchain.embeddings.cohere import CohereEmbeddings
 from langchain.embeddings.fake import FakeEmbeddings
@@ -24,6 +25,12 @@ from langchain.embeddings.localai import LocalAIEmbeddings
 from langchain.embeddings.openai import OpenAIEmbeddings
 
 
+@dataclass
+class Document:
+    page_content: str
+    record: AirbyteRecordMessage
+
+
 class Embedder(ABC):
     """
     Embedder is an abstract class that defines the interface for embedding text.
@@ -41,7 +48,7 @@ class Embedder(ABC):
         pass
 
     @abstractmethod
-    def embed_chunks(self, chunks: List[Chunk]) -> List[Optional[List[float]]]:
+    def embed_documents(self, documents: List[Document]) -> List[Optional[List[float]]]:
         """
         Embed the text of each chunk and return the resulting embedding vectors.
         If a chunk cannot be embedded or is configured to not be embedded, return None for that chunk.
@@ -72,7 +79,7 @@ class BaseOpenAIEmbedder(Embedder):
             return format_exception(e)
         return None
 
-    def embed_chunks(self, chunks: List[Chunk]) -> List[Optional[List[float]]]:
+    def embed_documents(self, documents: List[Document]) -> List[Optional[List[float]]]:
         """
         Embed the text of each chunk and return the resulting embedding vectors.
 
@@ -80,9 +87,9 @@ class BaseOpenAIEmbedder(Embedder):
         It's still possible to run into the rate limit between each embed call because the available token budget hasn't recovered between the calls,
         but the built-in retry mechanism of the OpenAI client handles that.
         """
-        # Each chunk can hold at most self.chunk_size tokens, so tokens-per-minute by maximum tokens per chunk is the number of chunks that can be embedded at once without exhausting the limit in a single request
+        # Each chunk can hold at most self.chunk_size tokens, so tokens-per-minute by maximum tokens per chunk is the number of documents that can be embedded at once without exhausting the limit in a single request
         embedding_batch_size = OPEN_AI_TOKEN_LIMIT // self.chunk_size
-        batches = create_chunks(chunks, batch_size=embedding_batch_size)
+        batches = create_chunks(documents, batch_size=embedding_batch_size)
         embeddings: List[Optional[List[float]]] = []
         for batch in batches:
             embeddings.extend(self.embeddings.embed_documents([chunk.page_content for chunk in batch]))
@@ -121,8 +128,8 @@ class CohereEmbedder(Embedder):
             return format_exception(e)
         return None
 
-    def embed_chunks(self, chunks: List[Chunk]) -> List[Optional[List[float]]]:
-        return cast(List[Optional[List[float]]], self.embeddings.embed_documents([chunk.page_content or "" for chunk in chunks]))
+    def embed_documents(self, documents: List[Document]) -> List[Optional[List[float]]]:
+        return cast(List[Optional[List[float]]], self.embeddings.embed_documents([document.page_content for document in documents]))
 
     @property
     def embedding_dimensions(self) -> int:
@@ -142,8 +149,8 @@ class FakeEmbedder(Embedder):
             return format_exception(e)
         return None
 
-    def embed_chunks(self, chunks: List[Chunk]) -> List[Optional[List[float]]]:
-        return cast(List[Optional[List[float]]], self.embeddings.embed_documents([chunk.page_content or "" for chunk in chunks]))
+    def embed_documents(self, documents: List[Document]) -> List[Optional[List[float]]]:
+        return cast(List[Optional[List[float]]], self.embeddings.embed_documents([document.page_content for document in documents]))
 
     @property
     def embedding_dimensions(self) -> int:
@@ -173,8 +180,8 @@ class OpenAICompatibleEmbedder(Embedder):
             return format_exception(e)
         return None
 
-    def embed_chunks(self, chunks: List[Chunk]) -> List[Optional[List[float]]]:
-        return cast(List[Optional[List[float]]], self.embeddings.embed_documents([chunk.page_content or "" for chunk in chunks]))
+    def embed_documents(self, documents: List[Document]) -> List[Optional[List[float]]]:
+        return cast(List[Optional[List[float]]], self.embeddings.embed_documents([document.page_content for document in documents]))
 
     @property
     def embedding_dimensions(self) -> int:
@@ -190,32 +197,32 @@ class FromFieldEmbedder(Embedder):
     def check(self) -> Optional[str]:
         return None
 
-    def embed_chunks(self, chunks: List[Chunk]) -> List[Optional[List[float]]]:
+    def embed_documents(self, documents: List[Document]) -> List[Optional[List[float]]]:
         """
         From each chunk, pull the embedding from the field specified in the config.
         Check that the field exists, is a list of numbers and is the correct size. If not, raise an AirbyteTracedException explaining the problem.
         """
         embeddings: List[Optional[List[float]]] = []
-        for chunk in chunks:
-            data = chunk.record.data
+        for document in documents:
+            data = document.record.data
             if self.config.field_name not in data:
                 raise AirbyteTracedException(
                     internal_message="Embedding vector field not found",
                     failure_type=FailureType.config_error,
-                    message=f"Record {str(data)[:250]}... in stream {chunk.record.stream}  does not contain embedding vector field {self.config.field_name}. Please check your embedding configuration, the embedding vector field has to be set correctly on every record.",
+                    message=f"Record {str(data)[:250]}... in stream {document.record.stream}  does not contain embedding vector field {self.config.field_name}. Please check your embedding configuration, the embedding vector field has to be set correctly on every record.",
                 )
             field = data[self.config.field_name]
             if not isinstance(field, list) or not all(isinstance(x, (int, float)) for x in field):
                 raise AirbyteTracedException(
                     internal_message="Embedding vector field not a list of numbers",
                     failure_type=FailureType.config_error,
-                    message=f"Record {str(data)[:250]}...  in stream {chunk.record.stream} does contain embedding vector field {self.config.field_name}, but it is not a list of numbers. Please check your embedding configuration, the embedding vector field has to be a list of numbers of length {self.config.dimensions} on every record.",
+                    message=f"Record {str(data)[:250]}...  in stream {document.record.stream} does contain embedding vector field {self.config.field_name}, but it is not a list of numbers. Please check your embedding configuration, the embedding vector field has to be a list of numbers of length {self.config.dimensions} on every record.",
                 )
             if len(field) != self.config.dimensions:
                 raise AirbyteTracedException(
                     internal_message="Embedding vector field has wrong length",
                     failure_type=FailureType.config_error,
-                    message=f"Record {str(data)[:250]}...  in stream {chunk.record.stream} does contain embedding vector field {self.config.field_name}, but it has length {len(field)} instead of the configured {self.config.dimensions}. Please check your embedding configuration, the embedding vector field has to be a list of numbers of length {self.config.dimensions} on every record.",
+                    message=f"Record {str(data)[:250]}...  in stream {document.record.stream} does contain embedding vector field {self.config.field_name}, but it has length {len(field)} instead of the configured {self.config.dimensions}. Please check your embedding configuration, the embedding vector field has to be a list of numbers of length {self.config.dimensions} on every record.",
                 )
             embeddings.append(field)
 

{airbyte-cdk-0.57.2 → airbyte-cdk-0.57.4}/airbyte_cdk/destinations/vector_db_based/writer.py

@@ -8,7 +8,7 @@ from typing import Dict, Iterable, List, Tuple
 
 from airbyte_cdk.destinations.vector_db_based.config import ProcessingConfigModel
 from airbyte_cdk.destinations.vector_db_based.document_processor import Chunk, DocumentProcessor
-from airbyte_cdk.destinations.vector_db_based.embedder import Embedder
+from airbyte_cdk.destinations.vector_db_based.embedder import Document, Embedder
 from airbyte_cdk.destinations.vector_db_based.indexer import Indexer
 from airbyte_cdk.models import AirbyteMessage, ConfiguredAirbyteCatalog, Type
 
@@ -16,14 +16,14 @@ from airbyte_cdk.models import AirbyteMessage, ConfiguredAirbyteCatalog, Type
 class Writer:
     """
     The Writer class is orchestrating the document processor, the embedder and the indexer:
-    * Incoming records are passed through the document processor to generate documents
-    * One the configured batch size is reached, the documents are passed to the embedder to generate embeddings
-    * The embedder embeds the documents
-    * The indexer deletes old documents by the associated record id before indexing the new ones
+    * Incoming records are passed through the document processor to generate chunks
+    * One the configured batch size is reached, the chunks are passed to the embedder to generate embeddings
+    * The embedder embeds the chunks
+    * The indexer deletes old chunks by the associated record id before indexing the new ones
 
     The destination connector is responsible to create a writer instance and pass the input messages iterable to the write method.
     The batch size can be configured by the destination connector to give the freedom of either letting the user configure it or hardcoding it to a sensible value depending on the destination.
-    The omit_raw_text parameter can be used to omit the raw text from the documents. This can be useful if the raw text is very large and not needed for the destination.
+    The omit_raw_text parameter can be used to omit the raw text from the chunks. This can be useful if the raw text is very large and not needed for the destination.
     """
 
     def __init__(
@@ -37,21 +37,29 @@ class Writer:
         self._init_batch()
 
     def _init_batch(self) -> None:
-        self.documents: Dict[Tuple[str, str], List[Chunk]] = defaultdict(list)
+        self.chunks: Dict[Tuple[str, str], List[Chunk]] = defaultdict(list)
         self.ids_to_delete: Dict[Tuple[str, str], List[str]] = defaultdict(list)
-        self.number_of_documents = 0
+        self.number_of_chunks = 0
+
+    def _convert_to_document(self, chunk: Chunk) -> Document:
+        """
+        Convert a chunk to a document for the embedder.
+        """
+        if chunk.page_content is None:
+            raise ValueError("Cannot embed a chunk without page content")
+        return Document(page_content=chunk.page_content, record=chunk.record)
 
     def _process_batch(self) -> None:
         for (namespace, stream), ids in self.ids_to_delete.items():
             self.indexer.delete(ids, namespace, stream)
 
-        for (namespace, stream), documents in self.documents.items():
-            embeddings = self.embedder.embed_chunks(documents)
-            for i, document in enumerate(documents):
+        for (namespace, stream), chunks in self.chunks.items():
+            embeddings = self.embedder.embed_documents([self._convert_to_document(chunk) for chunk in chunks])
+            for i, document in enumerate(chunks):
                 document.embedding = embeddings[i]
                 if self.omit_raw_text:
                     document.page_content = None
-            self.indexer.index(documents, namespace, stream)
+            self.indexer.index(chunks, namespace, stream)
 
         self._init_batch()
 
@@ -65,12 +73,12 @@ class Writer:
                 self._process_batch()
                 yield message
             elif message.type == Type.RECORD:
-                record_documents, record_id_to_delete = self.processor.process(message.record)
-                self.documents[(message.record.namespace, message.record.stream)].extend(record_documents)
+                record_chunks, record_id_to_delete = self.processor.process(message.record)
+                self.chunks[(message.record.namespace, message.record.stream)].extend(record_chunks)
                 if record_id_to_delete is not None:
                     self.ids_to_delete[(message.record.namespace, message.record.stream)].append(record_id_to_delete)
-                self.number_of_documents += len(record_documents)
-                if self.number_of_documents >= self.batch_size:
+                self.number_of_chunks += len(record_chunks)
+                if self.number_of_chunks >= self.batch_size:
                     self._process_batch()
 
         self._process_batch()

{airbyte-cdk-0.57.2 → airbyte-cdk-0.57.4}/airbyte_cdk/sources/file_based/file_types/unstructured_parser.py

@@ -3,6 +3,7 @@
 #
 import logging
 import traceback
+from datetime import datetime
 from io import BytesIO, IOBase
 from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Union
 
@@ -56,6 +57,8 @@ def user_error(e: Exception) -> bool:
     """
     Return True if this exception is caused by user error, False otherwise.
     """
+    if not isinstance(e, RecordParseError):
+        return False
     if not isinstance(e, requests.exceptions.RequestException):
         return False
     return bool(e.response and 400 <= e.response.status_code < 500)
@@ -164,10 +167,14 @@ class UnstructuredParser(FileTypeParser):
             return self._read_file_locally(file_handle, filetype, format.strategy, remote_file)
         elif format.processing.mode == "api":
             try:
-                result: str = self._read_file_remotely_with_retries(file_handle, format.processing, filetype, format.strategy)
+                result: str = self._read_file_remotely_with_retries(file_handle, format.processing, filetype, format.strategy, remote_file)
             except Exception as e:
-                # Re-throw as config error so the sync is stopped as problems with the external API need to be resolved by the user and are not considered part of the SLA.
+                # If a parser error happens during remotely processing the file, this means the file is corrupted. This case is handled by the parse_records method, so just rethrow.
+                #
+                # For other exceptions, re-throw as config error so the sync is stopped as problems with the external API need to be resolved by the user and are not considered part of the SLA.
                 # Once this parser leaves experimental stage, we should consider making this a system error instead for issues that might be transient.
+                if isinstance(e, RecordParseError):
+                    raise e
                 raise AirbyteTracedException.from_exception(e, failure_type=FailureType.config_error)
 
             return result
@@ -210,7 +217,13 @@ class UnstructuredParser(FileTypeParser):
             return False, "Base URL must start with https://"
 
         try:
-            self._read_file_remotely(BytesIO(b"# Airbyte source connection test"), format_config.processing, FileType.MD, "auto")
+            self._read_file_remotely(
+                BytesIO(b"# Airbyte source connection test"),
+                format_config.processing,
+                FileType.MD,
+                "auto",
+                RemoteFile(uri="test", last_modified=datetime.now()),
+            )
         except Exception:
             return False, "".join(traceback.format_exc())
 
@@ -218,14 +231,16 @@ class UnstructuredParser(FileTypeParser):
 
     @backoff.on_exception(backoff.expo, requests.exceptions.RequestException, max_tries=5, giveup=user_error)
     def _read_file_remotely_with_retries(
-        self, file_handle: IOBase, format: APIProcessingConfigModel, filetype: FileType, strategy: str
+        self, file_handle: IOBase, format: APIProcessingConfigModel, filetype: FileType, strategy: str, remote_file: RemoteFile
     ) -> str:
         """
         Read a file remotely, retrying up to 5 times if the error is not caused by user error. This is useful for transient network errors or the API server being overloaded temporarily.
         """
-        return self._read_file_remotely(file_handle, format, filetype, strategy)
+        return self._read_file_remotely(file_handle, format, filetype, strategy, remote_file)
 
-    def _read_file_remotely(self, file_handle: IOBase, format: APIProcessingConfigModel, filetype: FileType, strategy: str) -> str:
+    def _read_file_remotely(
+        self, file_handle: IOBase, format: APIProcessingConfigModel, filetype: FileType, strategy: str, remote_file: RemoteFile
+    ) -> str:
         headers = {"accept": "application/json", "unstructured-api-key": format.api_key}
 
         data = self._params_to_dict(format.parameters, strategy)
@@ -233,7 +248,13 @@ class UnstructuredParser(FileTypeParser):
         file_data = {"files": ("filename", file_handle, FILETYPE_TO_MIMETYPE[filetype])}
 
         response = requests.post(f"{format.api_url}/general/v0/general", headers=headers, data=data, files=file_data)
-        response.raise_for_status()
+
+        if response.status_code == 422:
+            # 422 means the file couldn't be processed, but the API is working. Treat this as a parsing error (passing an error record to the destination).
+            raise self._create_parse_error(remote_file, response.json())
+        else:
+            # Other error statuses are raised as requests exceptions (retry everything except user errors)
+            response.raise_for_status()
 
         json_response = response.json()
 

{airbyte-cdk-0.57.2 → airbyte-cdk-0.57.4}/airbyte_cdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 0.57.2
+Version: 0.57.4
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://github.com/airbytehq/airbyte
 Author: Airbyte

{airbyte-cdk-0.57.2 → airbyte-cdk-0.57.4}/setup.py

@@ -36,7 +36,7 @@ setup(
     name="airbyte-cdk",
     # The version of the airbyte-cdk package is used at runtime to validate manifests. That validation must be
     # updated if our semver format changes such as using release candidate versions.
-    version="0.57.2",
+    version="0.57.4",
     description="A framework for writing Airbyte Connectors.",
     long_description=README,
     long_description_content_type="text/markdown",

{airbyte-cdk-0.57.2 → airbyte-cdk-0.57.4}/unit_tests/sources/file_based/file_types/test_unstructured_parser.py

@@ -315,7 +315,7 @@ def test_check_config(requests_mock, format_config, raises_for_status, json_resp
 
 
 @pytest.mark.parametrize(
-    "filetype, format_config, raises_for_status, file_content, json_response, expected_requests, raises, expected_records",
+    "filetype, format_config, raises_for_status, file_content, json_response, expected_requests, raises, expected_records, http_status_code",
     [
         pytest.param(
             FileType.PDF,
@@ -332,6 +332,7 @@ def test_check_config(requests_mock, format_config, raises_for_status, json_resp
                     "_ab_source_file_parse_error": None
                 }
             ],
+            200,
             id="basic_request",
         ),
         pytest.param(
@@ -349,6 +350,7 @@ def test_check_config(requests_mock, format_config, raises_for_status, json_resp
                     "_ab_source_file_parse_error": None
                 }
             ],
+            200,
             id="request_with_params",
         ),
         pytest.param(
@@ -366,6 +368,7 @@ def test_check_config(requests_mock, format_config, raises_for_status, json_resp
                     "_ab_source_file_parse_error": None
                 }
             ],
+            200,
             id="handle_markdown_locally",
         ),
         pytest.param(
@@ -394,6 +397,7 @@ def test_check_config(requests_mock, format_config, raises_for_status, json_resp
             ],
             True,
             None,
+            200,
             id="retry_and_raise_on_api_error",
         ),
         pytest.param(
@@ -422,6 +426,7 @@ def test_check_config(requests_mock, format_config, raises_for_status, json_resp
                     "_ab_source_file_parse_error": None
                 }
             ],
+            200,
             id="retry_and_recover",
         ),
         pytest.param(
@@ -438,6 +443,7 @@ def test_check_config(requests_mock, format_config, raises_for_status, json_resp
             ],
             True,
             None,
+            200,
             id="no_retry_on_unexpected_error",
         ),
         pytest.param(
@@ -454,8 +460,29 @@ def test_check_config(requests_mock, format_config, raises_for_status, json_resp
             ],
             True,
             None,
+            400,
             id="no_retry_on_400_error",
         ),
+        pytest.param(
+            FileType.PDF,
+            UnstructuredFormat(skip_unprocessable_file_types=False, processing=APIProcessingConfigModel(mode="api", api_key="test")),
+            None,
+            "test",
+            [{"detail": "Something went wrong"}],
+            [
+                call("https://api.unstructured.io/general/v0/general", headers={"accept": "application/json", "unstructured-api-key": "test"}, data={"strategy": "auto"}, files={"files": ("filename", mock.ANY, "application/pdf")}),
+            ],
+            False,
+            [
+                {
+                    "content": None,
+                    "document_key": FILE_URI,
+                    "_ab_source_file_parse_error": "Error parsing record. This could be due to a mismatch between the config's file type and the actual file type, or because the file or record is not parseable. Contact Support if you need assistance.\nfilename=path/to/file.xyz message=[{'detail': 'Something went wrong'}]",
+                }
+            ],
+            422,
+            id="error_record_on_422_error",
+        ),
     ],
 )
 @patch("airbyte_cdk.sources.file_based.file_types.unstructured_parser.requests")
@@ -473,6 +500,7 @@ def test_parse_records_remotely(
     expected_requests,
     raises,
     expected_records,
+    http_status_code
 ):
     stream_reader = MagicMock()
     mock_open(stream_reader.open_file, read_data=bytes(str(file_content), "utf-8"))
@@ -484,6 +512,7 @@ def test_parse_records_remotely(
     mock_detect_filetype.return_value = filetype
     mock_response = MagicMock()
     mock_response.json.return_value = json_response
+    mock_response.status_code = http_status_code
     if raises_for_status:
         mock_response.raise_for_status.side_effect = raises_for_status
     requests_mock.post.return_value = mock_response