airbyte-cdk 0.51.11__tar.gz → 0.51.12__tar.gz

{airbyte-cdk-0.51.11 → airbyte-cdk-0.51.12}/PKG-INFO

@@ -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 → airbyte-cdk-0.51.12}/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-0.51.11 → airbyte-cdk-0.51.12}/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-0.51.12/airbyte_cdk/destinations/vector_db_based/embedder.py

@@ -0,0 +1,164 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+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
+
+
+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 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.
+    """
+
+    def __init__(self) -> None:
+        pass
+
+    @abstractmethod
+    def check(self) -> Optional[str]:
+        pass
+
+    @abstractmethod
+    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
+    @abstractmethod
+    def embedding_dimensions(self) -> int:
+        pass
+
+
+OPEN_AI_VECTOR_SIZE = 1536
+
+
+class OpenAIEmbedder(Embedder):
+    def __init__(self, config: OpenAIEmbeddingConfigModel):
+        super().__init__()
+        # Client is set internally
+        self.embeddings = OpenAIEmbeddings(openai_api_key=config.openai_key, chunk_size=8191, max_retries=15)  # type: ignore
+
+    def check(self) -> Optional[str]:
+        try:
+            self.embeddings.embed_query("test")
+        except Exception as e:
+            return format_exception(e)
+        return None
+
+    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:
+        # vector size produced by text-embedding-ada-002 model
+        return OPEN_AI_VECTOR_SIZE
+
+
+COHERE_VECTOR_SIZE = 1024
+
+
+class CohereEmbedder(Embedder):
+    def __init__(self, config: CohereEmbeddingConfigModel):
+        super().__init__()
+        # Client is set internally
+        self.embeddings = CohereEmbeddings(cohere_api_key=config.cohere_key, model="embed-english-light-v2.0")  # type: ignore
+
+    def check(self) -> Optional[str]:
+        try:
+            self.embeddings.embed_query("test")
+        except Exception as e:
+            return format_exception(e)
+        return None
+
+    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:
+        # vector size produced by text-embedding-ada-002 model
+        return COHERE_VECTOR_SIZE
+
+
+class FakeEmbedder(Embedder):
+    def __init__(self, config: FakeEmbeddingConfigModel):
+        super().__init__()
+        self.embeddings = FakeEmbeddings(size=OPEN_AI_VECTOR_SIZE)
+
+    def check(self) -> Optional[str]:
+        try:
+            self.embeddings.embed_query("test")
+        except Exception as e:
+            return format_exception(e)
+        return None
+
+    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-0.51.11 → airbyte-cdk-0.51.12}/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-0.51.11 → airbyte-cdk-0.51.12}/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-0.51.11 → airbyte-cdk-0.51.12}/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 → airbyte-cdk-0.51.12}/airbyte_cdk.egg-info/PKG-INFO

@@ -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 → airbyte-cdk-0.51.12}/setup.py

@@ -26,7 +26,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.51.11",
+    version="0.51.12",
     description="A framework for writing Airbyte Connectors.",
     long_description=README,
     long_description_content_type="text/markdown",

airbyte-cdk-0.51.11/airbyte_cdk/destinations/vector_db_based/embedder.py

@@ -1,109 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-
-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.utils import format_exception
-from langchain.embeddings.cohere import CohereEmbeddings
-from langchain.embeddings.fake import FakeEmbeddings
-from langchain.embeddings.openai import OpenAIEmbeddings
-
-
-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 CDK defines basic embedders that should be supported in each destination. It is possible to implement custom embedders for special destinations if needed.
-    """
-
-    def __init__(self) -> None:
-        pass
-
-    @abstractmethod
-    def check(self) -> Optional[str]:
-        pass
-
-    @abstractmethod
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
-        pass
-
-    @property
-    @abstractmethod
-    def embedding_dimensions(self) -> int:
-        pass
-
-
-OPEN_AI_VECTOR_SIZE = 1536
-
-
-class OpenAIEmbedder(Embedder):
-    def __init__(self, config: OpenAIEmbeddingConfigModel):
-        super().__init__()
-        # Client is set internally
-        self.embeddings = OpenAIEmbeddings(openai_api_key=config.openai_key, chunk_size=8191, max_retries=15)  # type: ignore
-
-    def check(self) -> Optional[str]:
-        try:
-            self.embeddings.embed_query("test")
-        except Exception as e:
-            return format_exception(e)
-        return None
-
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
-        return self.embeddings.embed_documents(texts)
-
-    @property
-    def embedding_dimensions(self) -> int:
-        # vector size produced by text-embedding-ada-002 model
-        return OPEN_AI_VECTOR_SIZE
-
-
-COHERE_VECTOR_SIZE = 1024
-
-
-class CohereEmbedder(Embedder):
-    def __init__(self, config: CohereEmbeddingConfigModel):
-        super().__init__()
-        # Client is set internally
-        self.embeddings = CohereEmbeddings(cohere_api_key=config.cohere_key, model="embed-english-light-v2.0")  # type: ignore
-
-    def check(self) -> Optional[str]:
-        try:
-            self.embeddings.embed_query("test")
-        except Exception as e:
-            return format_exception(e)
-        return None
-
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
-        return self.embeddings.embed_documents(texts)
-
-    @property
-    def embedding_dimensions(self) -> int:
-        # vector size produced by text-embedding-ada-002 model
-        return COHERE_VECTOR_SIZE
-
-
-class FakeEmbedder(Embedder):
-    def __init__(self, config: FakeEmbeddingConfigModel):
-        super().__init__()
-        self.embeddings = FakeEmbeddings(size=OPEN_AI_VECTOR_SIZE)
-
-    def check(self) -> Optional[str]:
-        try:
-            self.embeddings.embed_query("test")
-        except Exception as e:
-            return format_exception(e)
-        return None
-
-    def embed_texts(self, texts: List[str]) -> List[List[float]]:
-        return self.embeddings.embed_documents(texts)
-
-    @property
-    def embedding_dimensions(self) -> int:
-        # use same vector size as for OpenAI embeddings to keep it realistic
-        return OPEN_AI_VECTOR_SIZE