langchain/vectorstores/deeplake.py

"""Wrapper around Activeloop Deep Lake."""
from __future__ import annotations

import logging
import uuid
from typing import Any, Iterable, List, Optional, Sequence

import numpy as np

from langchain.docstore.document import Document
from langchain.embeddings.base import Embeddings
from langchain.vectorstores.base import VectorStore

logger = logging.getLogger()


def L2_search(
    query_embedding: np.ndarray, data_vectors: np.ndarray, k: int = 4
) -> list:
    """naive L2 search for nearest neighbors"""
    # Calculate the L2 distance between the query_vector and all data_vectors
    distances = np.linalg.norm(data_vectors - query_embedding, axis=1)

    # Sort the distances and return the indices of the k nearest vectors
    nearest_indices = np.argsort(distances)[:k]
    return nearest_indices.tolist()


class DeepLake(VectorStore):
    """Wrapper around Deep Lake, a data lake for deep learning applications.

    It not only stores embeddings, but also the original data and queries with
    version control automatically enabled.

    It is more than just a vector store. You can use the dataset to fine-tune
    your own LLM models or use it for other downstream tasks.

    We implement naive similiarity search, but it can be extended with Tensor
    Query Language (TQL for production use cases) over billion rows.

    To use, you should have the ``deeplake`` python package installed.

    Example:
        .. code-block:: python

                from langchain.vectorstores import DeepLake
                from langchain.embeddings.openai import OpenAIEmbeddings

                embeddings = OpenAIEmbeddings()
                vectorstore = DeepLake("langchain_store", embeddings.embed_query)
    """

    _LANGCHAIN_DEFAULT_DEEPLAKE_PATH = "mem://langchain"

    def __init__(
        self,
        dataset_path: str = _LANGCHAIN_DEFAULT_DEEPLAKE_PATH,
        token: Optional[str] = None,
        embedding_function: Optional[Embeddings] = None,
    ) -> None:
        """Initialize with Deep Lake client."""

        try:
            import deeplake
        except ImportError:
            raise ValueError(
                "Could not import deeplake python package. "
                "Please install it with `pip install deeplake`."
            )
        self._deeplake = deeplake

        if deeplake.exists(dataset_path, token=token):
            self.ds = deeplake.load(dataset_path, token=token)
            logger.warning(
                f"Deep Lake Dataset in {dataset_path} already exists, "
                f"loading from the storage"
            )
            self.ds.summary()
        else:
            self.ds = deeplake.empty(dataset_path, token=token, overwrite=True)
            with self.ds:
                self.ds.create_tensor("text", htype="text")
                self.ds.create_tensor("metadata", htype="json")
                self.ds.create_tensor("embedding", htype="generic")
                self.ds.create_tensor("ids", htype="text")

        self._embedding_function = embedding_function

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> List[str]:
        """Run more texts through the embeddings and add to the vectorstore.

        Args:
            texts (Iterable[str]): Texts to add to the vectorstore.
            metadatas (Optional[List[dict]], optional): Optional list of metadatas.
            ids (Optional[List[str]], optional): Optional list of IDs.

        Returns:
            List[str]: List of IDs of the added texts.
        """

        if ids is None:
            ids = [str(uuid.uuid1()) for _ in texts]

        text_list = list(texts)

        if self._embedding_function is None:
            embeddings: Sequence[Optional[List[float]]] = [None] * len(text_list)
        else:
            embeddings = self._embedding_function.embed_documents(text_list)

        if metadatas is None:
            metadatas_to_use: Sequence[Optional[dict]] = [None] * len(text_list)
        else:
            metadatas_to_use = metadatas

        elements = zip(text_list, embeddings, metadatas_to_use, ids)

        @self._deeplake.compute
        def ingest(sample_in: list, sample_out: list) -> None:
            s = {
                "text": sample_in[0],
                "embedding": sample_in[1],
                "metadata": sample_in[2],
                "ids": sample_in[3],
            }
            sample_out.append(s)

        ingest().eval(list(elements), self.ds)
        self.ds.commit()

        return ids

    def similarity_search(
        self, query: str, k: int = 4, **kwargs: Any
    ) -> List[Document]:
        """Return docs most similar to query."""
        if self._embedding_function is None:
            self.ds.summary()
            ds_view = self.ds.filter(lambda x: query in x["text"].data()["value"])
        else:
            query_emb = np.array(self._embedding_function.embed_query(query))
            embeddings = self.ds.embedding.numpy()
            indices = L2_search(query_emb, embeddings, k=k)
            ds_view = self.ds[indices]

        docs = [
            Document(
                page_content=el["text"].data()["value"],
                metadata=el["metadata"].data()["value"],
            )
            for el in ds_view
        ]
        return docs

    @classmethod
    def from_texts(
        cls,
        texts: List[str],
        embedding: Optional[Embeddings] = None,
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        dataset_path: str = _LANGCHAIN_DEFAULT_DEEPLAKE_PATH,
        **kwargs: Any,
    ) -> DeepLake:
        """Create a Deep Lake dataset from a raw documents.

        If a persist_directory is specified, the collection will be persisted there.
        Otherwise, the data will be ephemeral in-memory.

        Args:
            path (str, pathlib.Path): - The full path to the dataset. Can be:
                - a Deep Lake cloud path of the form ``hub://username/datasetname``.
                    To write to Deep Lake cloud datasets,
                    ensure that you are logged in to Deep Lake
                    (use 'activeloop login' from command line)
                - an s3 path of the form ``s3://bucketname/path/to/dataset``.
                    Credentials are required in either the environment or
                    passed to the creds argument.
                - a local file system path of the form ``./path/to/dataset`` or
                    ``~/path/to/dataset`` or ``path/to/dataset``.
                - a memory path of the form ``mem://path/to/dataset`` which doesn't
                    save the dataset but keeps it in memory instead.
                    Should be used only for testing as it does not persist.
            documents (List[Document]): List of documents to add.
            embedding (Optional[Embeddings]): Embedding function. Defaults to None.
            metadatas (Optional[List[dict]]): List of metadatas. Defaults to None.
            ids (Optional[List[str]]): List of document IDs. Defaults to None.

        Returns:
            DeepLake: Deep Lake dataset.
        """
        deeplake_dataset = cls(
            dataset_path=dataset_path,
            embedding_function=embedding,
        )
        deeplake_dataset.add_texts(texts=texts, metadatas=metadatas, ids=ids)
        return deeplake_dataset

    def delete_dataset(self) -> None:
        """Delete the collection."""
        self.ds.delete()

    def persist(self) -> None:
        """Persist the collection."""
        self.ds.flush()