OpenAIEmbeddingsDataSet

class OpenAIEmbeddingsDataSet(AbstractDataSet[None, OpenAIEmbeddings]):
    """OpenAI Embeddings dataset.
    Must be a dataset to access credentials at runtime.
    """

    def __init__(self, credentials: Dict[str, str], **kwargs):
        """Constructor.

        Args:
            credentials: must contain `openai_api_base` and `openai_api_key`.
            **kwargs: keyword arguments passed to the `OpenAIEmbeddings` class.
        """
        self.openai_api_base = credentials["openai_api_base"]
        self.openai_api_key = credentials["openai_api_key"]
        self.kwargs = kwargs

    def _describe(self) -> dict[str, Any]:
        return {**self.kwargs}

    def _save(self, data: None) -> NoReturn:
        raise DatasetError(f"{self.__class__.__name__} is a read only data set type")

    def _load(self) -> OpenAIEmbeddings:
        return OpenAIEmbeddings(
            openai_api_base=self.openai_api_base,
            openai_api_key=self.openai_api_key,
            **self.kwargs,
        )

from kedro.extras.datasets.pickle import PickleDataSet
from langchain.vectorstores import FAISS

class FAISSDataSet(PickleDataSet):
    """Saves and loads a FAISS vector store."""

    # TODO: find a better way to do this dataset.
    #   - Using anything but an API will also serialize the embedding model which will be too big.

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.openai_api_base = kwargs["credentials"]["openai_api_base"]
        self.openai_api_key = kwargs["credentials"]["openai_api_key"]

    def _load(self) -> FAISS:
        faiss = super()._load()

        # TODO: this assumes we're using an OpenAIEmbeddings embedding function.
        faiss.embedding_function.__self__.openai_api_base = self.openai_api_base
        faiss.embedding_function.__self__.openai_api_key = self.openai_api_key

        return faiss