airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/connector_builder/test_reader/types.py

@@ -0,0 +1,83 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+"""
+This module defines type aliases utilized in the Airbyte Connector Builder's test reader.
+These aliases streamline type-checking for heterogeneous message groups and schema outputs,
+ensuring consistency throughout the processing of stream data and associated messages.
+
+Type Aliases:
+    MESSAGE_GROUPS:
+        An iterable union of message-like objects which may include:
+            - StreamReadSlices: Represents slices used to read data from a stream.
+            - AirbyteControlMessage: Represents control commands used in the Airbyte protocol.
+            - AirbyteLogMessage: Represents log messages generated by the system.
+            - AirbyteTraceMessage: Represents trace messages typically used for debugging.
+            - AuxiliaryRequest: Represents any supplementary request issued during processing.
+
+    INFERRED_SCHEMA_OUTPUT_TYPE:
+        A tuple where:
+            - The first element is either an InferredSchema instance or None, denoting the inferred JSON schema.
+            - The second element is a list of LogMessage instances capturing logs produced during inference.
+
+    GROUPED_MESSAGES:
+        A tuple representing grouped messages divided as follows:
+            - A list of StreamReadSlices.
+            - A list of LogMessage instances.
+            - A list of AuxiliaryRequest instances.
+            - An optional AirbyteControlMessage that, if present, governs control flow in message processing.
+"""
+
+from typing import Any, Iterable, List
+
+from airbyte_cdk.connector_builder.models import (
+    AuxiliaryRequest,
+    HttpRequest,
+    HttpResponse,
+    LogMessage,
+    StreamReadSlices,
+)
+from airbyte_cdk.models import (
+    AirbyteControlMessage,
+    AirbyteLogMessage,
+    AirbyteTraceMessage,
+)
+from airbyte_cdk.utils.schema_inferrer import (
+    InferredSchema,
+)
+
+MESSAGE_GROUPS = Iterable[
+    StreamReadSlices
+    | AirbyteControlMessage
+    | AirbyteLogMessage
+    | AirbyteTraceMessage
+    | AuxiliaryRequest,
+]
+
+INFERRED_SCHEMA_OUTPUT_TYPE = tuple[
+    InferredSchema | None,
+    List[LogMessage],
+]
+
+GROUPED_MESSAGES = tuple[
+    List[StreamReadSlices],
+    List[LogMessage],
+    List[AuxiliaryRequest],
+    AirbyteControlMessage | None,
+]
+
+LOG_MESSAGES_OUTPUT_TYPE = tuple[
+    bool,
+    HttpRequest | None,
+    HttpResponse | None,
+    AuxiliaryRequest | None,
+    AirbyteLogMessage | None,
+]
+
+ASYNC_AUXILIARY_REQUEST_TYPES = [
+    "ASYNC_CREATE",
+    "ASYNC_POLL",
+    "ASYNC_ABORT",
+    "ASYNC_DELETE",
+]

airbyte_cdk/destinations/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+"""The destinations module provides classes for building destination connectors."""
+
+from .destination import Destination
+
+__all__ = [
+    "Destination",
+]

airbyte_cdk/destinations/destination.py

@@ -0,0 +1,154 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import argparse
+import io
+import logging
+import sys
+from abc import ABC, abstractmethod
+from typing import Any, Iterable, List, Mapping
+
+import orjson
+
+from airbyte_cdk.connector import Connector
+from airbyte_cdk.exception_handler import init_uncaught_exception_handler
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    AirbyteMessageSerializer,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
+    Type,
+)
+from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
+logger = logging.getLogger("airbyte")
+
+
+class Destination(Connector, ABC):
+    VALID_CMDS = {"spec", "check", "write"}
+
+    @abstractmethod
+    def write(
+        self,
+        config: Mapping[str, Any],
+        configured_catalog: ConfiguredAirbyteCatalog,
+        input_messages: Iterable[AirbyteMessage],
+    ) -> Iterable[AirbyteMessage]:
+        """Implement to define how the connector writes data to the destination"""
+
+    def _run_check(self, config: Mapping[str, Any]) -> AirbyteMessage:
+        check_result = self.check(logger, config)
+        return AirbyteMessage(type=Type.CONNECTION_STATUS, connectionStatus=check_result)
+
+    def _parse_input_stream(self, input_stream: io.TextIOWrapper) -> Iterable[AirbyteMessage]:
+        """Reads from stdin, converting to Airbyte messages"""
+        for line in input_stream:
+            try:
+                yield AirbyteMessageSerializer.load(orjson.loads(line))
+            except orjson.JSONDecodeError:
+                logger.info(
+                    f"ignoring input which can't be deserialized as Airbyte Message: {line}"
+                )
+
+    def _run_write(
+        self,
+        config: Mapping[str, Any],
+        configured_catalog_path: str,
+        input_stream: io.TextIOWrapper,
+    ) -> Iterable[AirbyteMessage]:
+        catalog = ConfiguredAirbyteCatalogSerializer.load(
+            orjson.loads(open(configured_catalog_path).read())
+        )
+        input_messages = self._parse_input_stream(input_stream)
+        logger.info("Begin writing to the destination...")
+        yield from self.write(
+            config=config, configured_catalog=catalog, input_messages=input_messages
+        )
+        logger.info("Writing complete.")
+
+    def parse_args(self, args: List[str]) -> argparse.Namespace:
+        """
+        :param args: commandline arguments
+        :return:
+        """
+
+        parent_parser = argparse.ArgumentParser(add_help=False)
+        main_parser = argparse.ArgumentParser()
+        subparsers = main_parser.add_subparsers(title="commands", dest="command")
+
+        # spec
+        subparsers.add_parser(
+            "spec", help="outputs the json configuration specification", parents=[parent_parser]
+        )
+
+        # check
+        check_parser = subparsers.add_parser(
+            "check", help="checks the config can be used to connect", parents=[parent_parser]
+        )
+        required_check_parser = check_parser.add_argument_group("required named arguments")
+        required_check_parser.add_argument(
+            "--config", type=str, required=True, help="path to the json configuration file"
+        )
+
+        # write
+        write_parser = subparsers.add_parser(
+            "write", help="Writes data to the destination", parents=[parent_parser]
+        )
+        write_required = write_parser.add_argument_group("required named arguments")
+        write_required.add_argument(
+            "--config", type=str, required=True, help="path to the JSON configuration file"
+        )
+        write_required.add_argument(
+            "--catalog", type=str, required=True, help="path to the configured catalog JSON file"
+        )
+
+        parsed_args = main_parser.parse_args(args)
+        cmd = parsed_args.command
+        if not cmd:
+            raise Exception("No command entered. ")
+        elif cmd not in ["spec", "check", "write"]:
+            # This is technically dead code since parse_args() would fail if this was the case
+            # But it's non-obvious enough to warrant placing it here anyways
+            raise Exception(f"Unknown command entered: {cmd}")
+
+        return parsed_args
+
+    def run_cmd(self, parsed_args: argparse.Namespace) -> Iterable[AirbyteMessage]:
+        cmd = parsed_args.command
+        if cmd not in self.VALID_CMDS:
+            raise Exception(f"Unrecognized command: {cmd}")
+
+        spec = self.spec(logger)
+        if cmd == "spec":
+            yield AirbyteMessage(type=Type.SPEC, spec=spec)
+            return
+        config = self.read_config(config_path=parsed_args.config)
+        if self.check_config_against_spec or cmd == "check":
+            try:
+                check_config_against_spec_or_exit(config, spec)
+            except AirbyteTracedException as traced_exc:
+                connection_status = traced_exc.as_connection_status_message()
+                if connection_status and cmd == "check":
+                    yield connection_status
+                    return
+                raise traced_exc
+
+        if cmd == "check":
+            yield self._run_check(config=config)
+        elif cmd == "write":
+            # Wrap in UTF-8 to override any other input encodings
+            wrapped_stdin = io.TextIOWrapper(sys.stdin.buffer, encoding="utf-8")
+            yield from self._run_write(
+                config=config,
+                configured_catalog_path=parsed_args.catalog,
+                input_stream=wrapped_stdin,
+            )
+
+    def run(self, args: List[str]) -> None:
+        init_uncaught_exception_handler(logger)
+        parsed_args = self.parse_args(args)
+        output_messages = self.run_cmd(parsed_args)
+        for message in output_messages:
+            print(orjson.dumps(AirbyteMessageSerializer.dump(message)).decode())

airbyte_cdk/destinations/vector_db_based/README.md

@@ -0,0 +1,37 @@
+# Vector DB based destinations
+
+## Note: All helpers in this directory are experimental and subject to change
+
+This directory contains several helpers that can be used to create a destination that processes and chunks records, embeds their text part and loads them into a vector database.
+The specific loading behavior is defined by the destination connector itself, but chunking and embedding behavior is handled by the helpers.
+
+To use these helpers, install the CDK with the `vector-db-based` extra:
+
+```bash
+pip install airbyte-cdk[vector-db-based]
+```
+
+The helpers can be used in the following way:
+
+- Add the config models to the spec of the connector
+- Implement the `Indexer` interface for your specific database
+- In the check implementation of the destination, initialize the indexer and the embedder and call `check` on them
+- In the write implementation of the destination, initialize the indexer, the embedder and pass them to a new instance of the writer. Then call the writers `write` method with the iterable for incoming messages
+
+If there are no connector-specific embedders, the `airbyte_cdk.destinations.vector_db_based.embedder.create_from_config` function can be used to get an embedder instance from the config.
+
+This is how the components interact:
+
+```text
+┌─────────────┐
+│MyDestination│
+└┬────────────┘
+┌▽───────────────────────────────┐
+│Writer                          │
+└┬─────────┬──────────┬──────────┘
+┌▽───────┐┌▽────────┐┌▽────────────────┐
+│Embedder││MyIndexer││DocumentProcessor│
+└────────┘└─────────┘└─────────────────┘
+```
+
+Normally, only the `MyDestination` class and the `MyIndexer` class has to be implemented specifically for the destination. The other classes are provided as is by the helpers.

airbyte_cdk/destinations/vector_db_based/__init__.py

@@ -0,0 +1,38 @@
+#
+# Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+#
+
+from .config import (
+    AzureOpenAIEmbeddingConfigModel,
+    CohereEmbeddingConfigModel,
+    FakeEmbeddingConfigModel,
+    FromFieldEmbeddingConfigModel,
+    OpenAICompatibleEmbeddingConfigModel,
+    OpenAIEmbeddingConfigModel,
+    ProcessingConfigModel,
+)
+from .document_processor import Chunk, DocumentProcessor
+from .embedder import CohereEmbedder, Embedder, FakeEmbedder, OpenAIEmbedder
+from .indexer import Indexer
+from .writer import Writer
+
+__all__ = [
+    "AzureOpenAIEmbedder",
+    "AzureOpenAIEmbeddingConfigModel",
+    "Chunk",
+    "CohereEmbedder",
+    "CohereEmbeddingConfigModel",
+    "DocumentProcessor",
+    "Embedder",
+    "FakeEmbedder",
+    "FakeEmbeddingConfigModel",
+    "FromFieldEmbedder",
+    "FromFieldEmbeddingConfigModel",
+    "Indexer",
+    "OpenAICompatibleEmbedder",
+    "OpenAICompatibleEmbeddingConfigModel",
+    "OpenAIEmbedder",
+    "OpenAIEmbeddingConfigModel",
+    "ProcessingConfigModel",
+    "Writer",
+]

airbyte_cdk/destinations/vector_db_based/config.py

@@ -0,0 +1,298 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import dpath
+from pydantic.v1 import BaseModel, Field
+
+from airbyte_cdk.utils.oneof_option_config import OneOfOptionConfig
+from airbyte_cdk.utils.spec_schema_transformations import resolve_refs
+
+
+class SeparatorSplitterConfigModel(BaseModel):
+    mode: Literal["separator"] = Field("separator", const=True)
+    separators: List[str] = Field(
+        default=['"\\n\\n"', '"\\n"', '" "', '""'],
+        title="Separators",
+        description='List of separator strings to split text fields by. The separator itself needs to be wrapped in double quotes, e.g. to split by the dot character, use ".". To split by a newline, use "\\n".',
+    )
+    keep_separator: bool = Field(
+        default=False,
+        title="Keep separator",
+        description="Whether to keep the separator in the resulting chunks",
+    )
+
+    class Config(OneOfOptionConfig):
+        title = "By Separator"
+        description = "Split the text by the list of separators until the chunk size is reached, using the earlier mentioned separators where possible. This is useful for splitting text fields by paragraphs, sentences, words, etc."
+        discriminator = "mode"
+
+
+class MarkdownHeaderSplitterConfigModel(BaseModel):
+    mode: Literal["markdown"] = Field("markdown", const=True)
+    split_level: int = Field(
+        default=1,
+        title="Split level",
+        description="Level of markdown headers to split text fields by. Headings down to the specified level will be used as split points",
+        le=6,
+        ge=1,
+    )
+
+    class Config(OneOfOptionConfig):
+        title = "By Markdown header"
+        description = "Split the text by Markdown headers down to the specified header level. If the chunk size fits multiple sections, they will be combined into a single chunk."
+        discriminator = "mode"
+
+
+class CodeSplitterConfigModel(BaseModel):
+    mode: Literal["code"] = Field("code", const=True)
+    language: str = Field(
+        title="Language",
+        description="Split code in suitable places based on the programming language",
+        enum=[
+            "cpp",
+            "go",
+            "java",
+            "js",
+            "php",
+            "proto",
+            "python",
+            "rst",
+            "ruby",
+            "rust",
+            "scala",
+            "swift",
+            "markdown",
+            "latex",
+            "html",
+            "sol",
+        ],
+    )
+
+    class Config(OneOfOptionConfig):
+        title = "By Programming Language"
+        description = "Split the text by suitable delimiters based on the programming language. This is useful for splitting code into chunks."
+        discriminator = "mode"
+
+
+TextSplitterConfigModel = Union[
+    SeparatorSplitterConfigModel, MarkdownHeaderSplitterConfigModel, CodeSplitterConfigModel
+]
+
+
+class FieldNameMappingConfigModel(BaseModel):
+    from_field: str = Field(title="From field name", description="The field name in the source")
+    to_field: str = Field(
+        title="To field name", description="The field name to use in the destination"
+    )
+
+
+class ProcessingConfigModel(BaseModel):
+    chunk_size: int = Field(
+        ...,
+        title="Chunk size",
+        maximum=8191,
+        minimum=1,
+        description="Size of chunks in tokens to store in vector store (make sure it is not too big for the context if your LLM)",
+    )
+    chunk_overlap: int = Field(
+        title="Chunk overlap",
+        description="Size of overlap between chunks in tokens to store in vector store to better capture relevant context",
+        default=0,
+    )
+    text_fields: Optional[List[str]] = Field(
+        default=[],
+        title="Text fields to embed",
+        description="List of fields in the record that should be used to calculate the embedding. The field list is applied to all streams in the same way and non-existing fields are ignored. If none are defined, all fields are considered text fields. When specifying text fields, you can access nested fields in the record by using dot notation, e.g. `user.name` will access the `name` field in the `user` object. It's also possible to use wildcards to access all fields in an object, e.g. `users.*.name` will access all `names` fields in all entries of the `users` array.",
+        always_show=True,
+        examples=["text", "user.name", "users.*.name"],
+    )
+    metadata_fields: Optional[List[str]] = Field(
+        default=[],
+        title="Fields to store as metadata",
+        description="List of fields in the record that should be stored as metadata. The field list is applied to all streams in the same way and non-existing fields are ignored. If none are defined, all fields are considered metadata fields. When specifying text fields, you can access nested fields in the record by using dot notation, e.g. `user.name` will access the `name` field in the `user` object. It's also possible to use wildcards to access all fields in an object, e.g. `users.*.name` will access all `names` fields in all entries of the `users` array. When specifying nested paths, all matching values are flattened into an array set to a field named by the path.",
+        always_show=True,
+        examples=["age", "user", "user.name"],
+    )
+    text_splitter: TextSplitterConfigModel = Field(
+        default=None,
+        title="Text splitter",
+        discriminator="mode",
+        type="object",
+        description="Split text fields into chunks based on the specified method.",
+    )
+    field_name_mappings: Optional[List[FieldNameMappingConfigModel]] = Field(
+        default=[],
+        title="Field name mappings",
+        description="List of fields to rename. Not applicable for nested fields, but can be used to rename fields already flattened via dot notation.",
+    )
+
+    class Config:
+        schema_extra = {"group": "processing"}
+
+
+class OpenAIEmbeddingConfigModel(BaseModel):
+    mode: Literal["openai"] = Field("openai", const=True)
+    openai_key: str = Field(..., title="OpenAI API key", airbyte_secret=True)
+
+    class Config(OneOfOptionConfig):
+        title = "OpenAI"
+        description = "Use the OpenAI API to embed text. This option is using the text-embedding-ada-002 model with 1536 embedding dimensions."
+        discriminator = "mode"
+
+
+class OpenAICompatibleEmbeddingConfigModel(BaseModel):
+    mode: Literal["openai_compatible"] = Field("openai_compatible", const=True)
+    api_key: str = Field(title="API key", default="", airbyte_secret=True)
+    base_url: str = Field(
+        ...,
+        title="Base URL",
+        description="The base URL for your OpenAI-compatible service",
+        examples=["https://your-service-name.com"],
+    )
+    model_name: str = Field(
+        title="Model name",
+        description="The name of the model to use for embedding",
+        default="text-embedding-ada-002",
+        examples=["text-embedding-ada-002"],
+    )
+    dimensions: int = Field(
+        title="Embedding dimensions",
+        description="The number of dimensions the embedding model is generating",
+        examples=[1536, 384],
+    )
+
+    class Config(OneOfOptionConfig):
+        title = "OpenAI-compatible"
+        description = "Use a service that's compatible with the OpenAI API to embed text."
+        discriminator = "mode"
+
+
+class AzureOpenAIEmbeddingConfigModel(BaseModel):
+    mode: Literal["azure_openai"] = Field("azure_openai", const=True)
+    openai_key: str = Field(
+        ...,
+        title="Azure OpenAI API key",
+        airbyte_secret=True,
+        description="The API key for your Azure OpenAI resource.  You can find this in the Azure portal under your Azure OpenAI resource",
+    )
+    api_base: str = Field(
+        ...,
+        title="Resource base URL",
+        description="The base URL for your Azure OpenAI resource.  You can find this in the Azure portal under your Azure OpenAI resource",
+        examples=["https://your-resource-name.openai.azure.com"],
+    )
+    deployment: str = Field(
+        ...,
+        title="Deployment",
+        description="The deployment for your Azure OpenAI resource.  You can find this in the Azure portal under your Azure OpenAI resource",
+        examples=["your-resource-name"],
+    )
+
+    class Config(OneOfOptionConfig):
+        title = "Azure OpenAI"
+        description = "Use the Azure-hosted OpenAI API to embed text. This option is using the text-embedding-ada-002 model with 1536 embedding dimensions."
+        discriminator = "mode"
+
+
+class FakeEmbeddingConfigModel(BaseModel):
+    mode: Literal["fake"] = Field("fake", const=True)
+
+    class Config(OneOfOptionConfig):
+        title = "Fake"
+        description = "Use a fake embedding made out of random vectors with 1536 embedding dimensions. This is useful for testing the data pipeline without incurring any costs."
+        discriminator = "mode"
+
+
+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(OneOfOptionConfig):
+        title = "From Field"
+        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."
+        discriminator = "mode"
+
+
+class CohereEmbeddingConfigModel(BaseModel):
+    mode: Literal["cohere"] = Field("cohere", const=True)
+    cohere_key: str = Field(..., title="Cohere API key", airbyte_secret=True)
+
+    class Config(OneOfOptionConfig):
+        title = "Cohere"
+        description = "Use the Cohere API to embed text."
+        discriminator = "mode"
+
+
+class VectorDBConfigModel(BaseModel):
+    """
+    The configuration model for the Vector DB based destinations. This model is used to generate the UI for the destination configuration,
+    as well as to provide type safety for the configuration passed to the destination.
+
+    The configuration model is composed of four parts:
+    * Processing configuration
+    * Embedding configuration
+    * Indexing configuration
+    * Advanced configuration
+
+    Processing, embedding and advanced configuration are provided by this base class, while the indexing configuration is provided by the destination connector in the sub class.
+    """
+
+    embedding: Union[
+        OpenAIEmbeddingConfigModel,
+        CohereEmbeddingConfigModel,
+        FakeEmbeddingConfigModel,
+        AzureOpenAIEmbeddingConfigModel,
+        OpenAICompatibleEmbeddingConfigModel,
+    ] = Field(
+        ...,
+        title="Embedding",
+        description="Embedding configuration",
+        discriminator="mode",
+        group="embedding",
+        type="object",
+    )
+    processing: ProcessingConfigModel
+    omit_raw_text: bool = Field(
+        default=False,
+        title="Do not store raw text",
+        group="advanced",
+        description="Do not store the text that gets embedded along with the vector and the metadata in the destination. If set to true, only the vector and the metadata will be stored - in this case raw text for LLM use cases needs to be retrieved from another source.",
+    )
+
+    class Config:
+        title = "Destination Config"
+        schema_extra = {
+            "groups": [
+                {"id": "processing", "title": "Processing"},
+                {"id": "embedding", "title": "Embedding"},
+                {"id": "indexing", "title": "Indexing"},
+                {"id": "advanced", "title": "Advanced"},
+            ]
+        }
+
+    @staticmethod
+    def remove_discriminator(schema: Dict[str, Any]) -> None:
+        """pydantic adds "discriminator" to the schema for oneOfs, which is not treated right by the platform as we inline all references"""
+        dpath.delete(schema, "properties/**/discriminator")
+
+    @classmethod
+    def schema(cls, by_alias: bool = True, ref_template: str = "") -> Dict[str, Any]:
+        """we're overriding the schema classmethod to enable some post-processing"""
+        schema: Dict[str, Any] = super().schema()
+        schema = resolve_refs(schema)
+        cls.remove_discriminator(schema)
+        return schema