airbyte-cdk 0.51.11__py3-none-any.whl → 0.51.12__py3-none-any.whl

airbyte_cdk/destinations/vector_db_based/config.py

@@ -59,6 +59,22 @@ class FakeEmbeddingConfigModel(BaseModel):
         }
 
 
+class FromFieldEmbeddingConfigModel(BaseModel):
+    mode: Literal["from_field"] = Field("from_field", const=True)
+    field_name: str = Field(
+        ..., title="Field name", description="Name of the field in the record that contains the embedding", examples=["embedding", "vector"]
+    )
+    dimensions: int = Field(
+        ..., title="Embedding dimensions", description="The number of dimensions the embedding model is generating", examples=[1536, 384]
+    )
+
+    class Config:
+        title = "From Field"
+        schema_extra = {
+            "description": "Use a field in the record as the embedding. This is useful if you already have an embedding for your data and want to store it in the vector store."
+        }
+
+
 class CohereEmbeddingConfigModel(BaseModel):
     mode: Literal["cohere"] = Field("cohere", const=True)
     cohere_key: str = Field(..., title="Cohere API key", airbyte_secret=True)

airbyte_cdk/destinations/vector_db_based/document_processor.py

@@ -9,6 +9,7 @@ from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
 import dpath.util
 from airbyte_cdk.destinations.vector_db_based.config import ProcessingConfigModel
 from airbyte_cdk.models import AirbyteRecordMessage, AirbyteStream, ConfiguredAirbyteCatalog, ConfiguredAirbyteStream, DestinationSyncMode
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException, FailureType
 from langchain.document_loaders.base import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.utils import stringify_dict
@@ -21,8 +22,8 @@ METADATA_RECORD_ID_FIELD = "_ab_record_id"
 class Chunk:
     page_content: str
     metadata: Dict[str, Any]
-    stream: str
-    namespace: Optional[str] = None
+    record: AirbyteRecordMessage
+    embedding: Optional[List[float]] = None
 
 
 class DocumentProcessor:
@@ -66,11 +67,14 @@ class DocumentProcessor:
         """
         doc = self._generate_document(record)
         if doc is None:
-            raise ValueError(f"Record {str(record.data)[:250]}... does not contain any text fields.")
-        chunks = [
-            Chunk(
-                page_content=chunk_document.page_content, metadata=chunk_document.metadata, stream=record.stream, namespace=record.namespace
+            text_fields = ", ".join(self.text_fields) if self.text_fields else "all fields"
+            raise AirbyteTracedException(
+                internal_message="No text fields found in record",
+                message=f"Record {str(record.data)[:250]}... does not contain any of the configured text fields: {text_fields}. Please check your processing configuration, there has to be at least one text field set in each record.",
+                failure_type=FailureType.config_error,
             )
+        chunks = [
+            Chunk(page_content=chunk_document.page_content, metadata=chunk_document.metadata, record=record)
             for chunk_document in self._split_document(doc)
         ]
         id_to_delete = doc.metadata[METADATA_RECORD_ID_FIELD] if METADATA_RECORD_ID_FIELD in doc.metadata else None

airbyte_cdk/destinations/vector_db_based/embedder.py

@@ -5,8 +5,15 @@
 from abc import ABC, abstractmethod
 from typing import List, Optional
 
-from airbyte_cdk.destinations.vector_db_based.config import CohereEmbeddingConfigModel, FakeEmbeddingConfigModel, OpenAIEmbeddingConfigModel
+from airbyte_cdk.destinations.vector_db_based.config import (
+    CohereEmbeddingConfigModel,
+    FakeEmbeddingConfigModel,
+    FromFieldEmbeddingConfigModel,
+    OpenAIEmbeddingConfigModel,
+)
+from airbyte_cdk.destinations.vector_db_based.document_processor import Chunk
 from airbyte_cdk.destinations.vector_db_based.utils import format_exception
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException, FailureType
 from langchain.embeddings.cohere import CohereEmbeddings
 from langchain.embeddings.fake import FakeEmbeddings
 from langchain.embeddings.openai import OpenAIEmbeddings
@@ -17,7 +24,7 @@ class Embedder(ABC):
     Embedder is an abstract class that defines the interface for embedding text.
 
     The Indexer class uses the Embedder class to internally embed text - each indexer is responsible to pass the text of all documents to the embedder and store the resulting embeddings in the destination.
-    The destination connector is responsible to create an embedder instance and pass it to the indexer.
+    The destination connector is responsible to create an embedder instance and pass it to the writer.
     The CDK defines basic embedders that should be supported in each destination. It is possible to implement custom embedders for special destinations if needed.
     """
 
@@ -29,7 +36,11 @@ class Embedder(ABC):
         pass
 
     @abstractmethod
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
+    def embed_chunks(self, chunks: List[Chunk]) -> 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.
+        """
         pass
 
     @property
@@ -54,8 +65,8 @@ class OpenAIEmbedder(Embedder):
             return format_exception(e)
         return None
 
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
-        return self.embeddings.embed_documents(texts)
+    def embed_chunks(self, chunks: List[Chunk]) -> List[List[float]]:
+        return self.embeddings.embed_documents([chunk.page_content for chunk in chunks])
 
     @property
     def embedding_dimensions(self) -> int:
@@ -79,8 +90,8 @@ class CohereEmbedder(Embedder):
             return format_exception(e)
         return None
 
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
-        return self.embeddings.embed_documents(texts)
+    def embed_chunks(self, chunks: List[Chunk]) -> List[List[float]]:
+        return self.embeddings.embed_documents([chunk.page_content for chunk in chunks])
 
     @property
     def embedding_dimensions(self) -> int:
@@ -100,10 +111,54 @@ class FakeEmbedder(Embedder):
             return format_exception(e)
         return None
 
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
-        return self.embeddings.embed_documents(texts)
+    def embed_chunks(self, chunks: List[Chunk]) -> List[List[float]]:
+        return self.embeddings.embed_documents([chunk.page_content for chunk in chunks])
 
     @property
     def embedding_dimensions(self) -> int:
         # use same vector size as for OpenAI embeddings to keep it realistic
         return OPEN_AI_VECTOR_SIZE
+
+
+class FromFieldEmbedder(Embedder):
+    def __init__(self, config: FromFieldEmbeddingConfigModel):
+        super().__init__()
+        self.config = config
+
+    def check(self) -> Optional[str]:
+        return None
+
+    def embed_chunks(self, chunks: List[Chunk]) -> List[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 = []
+        for chunk in chunks:
+            data = chunk.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.",
+                )
+            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.",
+                )
+            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.",
+                )
+            embeddings.append(field)
+
+        return embeddings
+
+    @property
+    def embedding_dimensions(self) -> int:
+        return self.config.dimensions

airbyte_cdk/destinations/vector_db_based/indexer.py

@@ -7,7 +7,6 @@ from abc import ABC, abstractmethod
 from typing import Any, Generator, Iterable, List, Optional, Tuple, TypeVar
 
 from airbyte_cdk.destinations.vector_db_based.document_processor import Chunk
-from airbyte_cdk.destinations.vector_db_based.embedder import Embedder
 from airbyte_cdk.models import AirbyteMessage, ConfiguredAirbyteCatalog
 
 
@@ -19,9 +18,8 @@ class Indexer(ABC):
     In a destination connector, implement a custom indexer by extending this class and implementing the abstract methods.
     """
 
-    def __init__(self, config: Any, embedder: Embedder):
+    def __init__(self, config: Any):
         self.config = config
-        self.embedder = embedder
         pass
 
     def pre_sync(self, catalog: ConfiguredAirbyteCatalog) -> None:

airbyte_cdk/destinations/vector_db_based/test_utils.py

@@ -48,6 +48,6 @@ class BaseIntegrationTest(unittest.TestCase):
             type=Type.RECORD, record=AirbyteRecordMessage(stream=stream, data={"str_col": str_value, "int_col": int_value}, emitted_at=0)
         )
 
-    def setUp(self):
+    def setUp(self) -> None:
         with open("secrets/config.json", "r") as f:
             self.config = json.loads(f.read())

airbyte_cdk/destinations/vector_db_based/writer.py

@@ -8,24 +8,27 @@ from typing import Iterable, List
 from airbyte_cdk.destinations.vector_db_based.batcher import Batcher
 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.indexer import Indexer
 from airbyte_cdk.models import AirbyteMessage, AirbyteRecordMessage, ConfiguredAirbyteCatalog, Type
 
 
 class Writer:
     """
-    The Writer class is orchestrating the document processor, the batcher and the indexer:
+    The Writer class is orchestrating the document processor, the batcher, the embedder and the indexer:
     * Incoming records are collected using the batcher
     * The document processor generates documents from all records in the batch
-    * The indexer indexes the resulting documents in the destination
+    * The embedder embeds the documents
+    * The indexer indexes the resulting documents and their embeddings in the destination
 
     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.
     """
 
-    def __init__(self, processing_config: ProcessingConfigModel, indexer: Indexer, batch_size: int) -> None:
+    def __init__(self, processing_config: ProcessingConfigModel, indexer: Indexer, embedder: Embedder, batch_size: int) -> None:
         self.processing_config = processing_config
         self.indexer = indexer
+        self.embedder = embedder
         self.batcher = Batcher(batch_size, lambda batch: self._process_batch(batch))
 
     def _process_batch(self, batch: List[AirbyteRecordMessage]) -> None:
@@ -36,6 +39,9 @@ class Writer:
             documents.extend(record_documents)
             if record_id_to_delete is not None:
                 ids_to_delete.append(record_id_to_delete)
+        embeddings = self.embedder.embed_chunks(documents)
+        for i, document in enumerate(documents):
+            document.embedding = embeddings[i]
         self.indexer.index(documents, ids_to_delete)
 
     def write(self, configured_catalog: ConfiguredAirbyteCatalog, input_messages: Iterable[AirbyteMessage]) -> Iterable[AirbyteMessage]:

{airbyte_cdk-0.51.11.dist-info → airbyte_cdk-0.51.12.dist-info}/METADATA

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

{airbyte_cdk-0.51.11.dist-info → airbyte_cdk-0.51.12.dist-info}/RECORD

@@ -14,13 +14,13 @@ airbyte_cdk/destinations/__init__.py,sha256=0Uxmz3iBAyZJdk_bqUVt2pb0UwRTpFjTnFE6
 airbyte_cdk/destinations/destination.py,sha256=_tIMnKcRQbtIsjVvNOVjfbIxgCNLuBXQwQj8MyVm3BI,5420
 airbyte_cdk/destinations/vector_db_based/__init__.py,sha256=z5Pqxxt3v-JCcJQ6sK4tAz5sg1FB-3wTCd2p85MhFzc,711
 airbyte_cdk/destinations/vector_db_based/batcher.py,sha256=U2RI0CACZ1WhJIdkC5oPlwZ90OZB40kyFCR5I7StqZw,1160
-airbyte_cdk/destinations/vector_db_based/config.py,sha256=JrK2yuxWGaBa2PBP8WUX8WMG2ah6ah-Z75nWCeoGm_A,3407
-airbyte_cdk/destinations/vector_db_based/document_processor.py,sha256=D4L2NiMLv-Wy2FuYjluk-tuHljmSWMYGLHPNwoMi_0c,5716
-airbyte_cdk/destinations/vector_db_based/embedder.py,sha256=mB-VOVB3-DCKEtKCNhcuLH9OFFcyuP3yNCUf37OtV1M,3640
-airbyte_cdk/destinations/vector_db_based/indexer.py,sha256=WFSPGJsvCRSqdZkgeM2RCSD8VPgHOnJmKbXSNK5XZos,2451
-airbyte_cdk/destinations/vector_db_based/test_utils.py,sha256=Ldf-nJYCCzemP0bkjMQiRiApNh252cBs8PHKlFabG6o,1799
+airbyte_cdk/destinations/vector_db_based/config.py,sha256=xv5-IhPG_eKdRxstYmaFBUrYDECevE64OVRyUBZAJJw,4132
+airbyte_cdk/destinations/vector_db_based/document_processor.py,sha256=KHvCSjt6amwpIYxK42OuT1Vh-RCA5A3vEBfAmowXpZI,6161
+airbyte_cdk/destinations/vector_db_based/embedder.py,sha256=0YLm5wmqiwCyUD_GWzqetWclzzKsADOfjXu0jMhQS1Y,6837
+airbyte_cdk/destinations/vector_db_based/indexer.py,sha256=DMic7D7ie4gGQ-yOgGXGYjBsY8H7X5O5Tz_sCr0ajBU,2327
+airbyte_cdk/destinations/vector_db_based/test_utils.py,sha256=8d1Smk4jQRKtDfloXfEq12T-BU8ByyzzSBwAlchsU4A,1807
 airbyte_cdk/destinations/vector_db_based/utils.py,sha256=ngJ6hc9mmzgAEEBd9nuoRcPPFUKijv2CA6zZYUVRm54,240
-airbyte_cdk/destinations/vector_db_based/writer.py,sha256=uf3QLEKwaJsuGuhbfaJ4qbvr8EXc4NXr4J3JERP5dtI,2756
+airbyte_cdk/destinations/vector_db_based/writer.py,sha256=zSVizVPupTjdF_dwniIU0RYnTZ9TMkizOK48tDNPxxk,3110
 airbyte_cdk/models/__init__.py,sha256=rDARocDgxf4_qI66Bm6dHTBoecbWguTClGVBmOBiI2o,1674
 airbyte_cdk/models/airbyte_protocol.py,sha256=DoJvnmGM3xMAZFTwA6_RGMiKSFqfE3ib_Ru0KJ65Ag4,100
 airbyte_cdk/models/well_known_types.py,sha256=KKfNbow2gdLoC1Z4hcXy_JR8m_acsB2ol7gQuEgjobw,117
@@ -365,8 +365,8 @@ unit_tests/utils/test_schema_inferrer.py,sha256=Z2jHBZ540wnYkylIdV_2xr75Vtwlxuyg
 unit_tests/utils/test_secret_utils.py,sha256=XKe0f1RHYii8iwE6ATmBr5JGDI1pzzrnZUGdUSMJQP4,4886
 unit_tests/utils/test_stream_status_utils.py,sha256=NpV155JMXA6CG-2Zvofa14lItobyh3Onttc59X4m5DI,3382
 unit_tests/utils/test_traced_exception.py,sha256=bDFP5zMBizFenz6V2WvEZTRCKGB5ijh3DBezjbfoYIs,4198
-airbyte_cdk-0.51.11.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
-airbyte_cdk-0.51.11.dist-info/METADATA,sha256=2yDh-4NTXY4ALKd3svXb7IoBxTefflum6uw2szW9cjw,9895
-airbyte_cdk-0.51.11.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-airbyte_cdk-0.51.11.dist-info/top_level.txt,sha256=edvsDKTnE6sD2wfCUaeTfKf5gQIL6CPVMwVL2sWZzqo,51
-airbyte_cdk-0.51.11.dist-info/RECORD,,
+airbyte_cdk-0.51.12.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
+airbyte_cdk-0.51.12.dist-info/METADATA,sha256=kP39_c0A5hJ-e8yU-oZ-zAbknbhsFKaz7I11AoLyh5o,9895
+airbyte_cdk-0.51.12.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+airbyte_cdk-0.51.12.dist-info/top_level.txt,sha256=edvsDKTnE6sD2wfCUaeTfKf5gQIL6CPVMwVL2sWZzqo,51
+airbyte_cdk-0.51.12.dist-info/RECORD,,