Toggle Menu
Prev Up Next

# Source code for langchain_community.vectorstores.rocksetdb

    from __future__ import annotations

    import logging
    from copy import deepcopy
    from enum import Enum
    from typing import Any, Iterable, List, Optional, Tuple

    import numpy as np
    from langchain_core.documents import Document
    from langchain_core.embeddings import Embeddings
    from langchain_core.runnables import run_in_executor
    from langchain_core.vectorstores import VectorStore

    from langchain_community.vectorstores.utils import maximal_marginal_relevance

    logger = logging.getLogger(__name__)

    class Rockset(VectorStore):
        """`Rockset` vector store.

        To use, you should have the `rockset` python package installed. Note that to use
        this, the collection being used must already exist in your Rockset instance.
        You must also ensure you use a Rockset ingest transformation to apply
        `VECTOR_ENFORCE` on the column being used to store `embedding_key` in the
        collection.
        See: https://rockset.com/blog/introducing-vector-search-on-rockset/ for more details

        Everything below assumes `commons` Rockset workspace.

        Example:
            .. code-block:: python

                from langchain_community.vectorstores import Rockset
                from langchain_community.embeddings.openai import OpenAIEmbeddings
                import rockset

                # Make sure you use the right host (region) for your Rockset instance
                # and APIKEY has both read-write access to your collection.

                rs = rockset.RocksetClient(host=rockset.Regions.use1a1, api_key="***")
                collection_name = "langchain_demo"
                embeddings = OpenAIEmbeddings()
                vectorstore = Rockset(rs, collection_name, embeddings,
                    "description", "description_embedding")

        """

        def __init__(
            self,
            client: Any,
            embeddings: Embeddings,
            collection_name: str,
            text_key: str,
            embedding_key: str,
            workspace: str = "commons",
        ):
            """Initialize with Rockset client.
            Args:
                client: Rockset client object
                collection: Rockset collection to insert docs / query
                embeddings: Langchain Embeddings object to use to generate
                            embedding for given text.
                text_key: column in Rockset collection to use to store the text
                embedding_key: column in Rockset collection to use to store the embedding.
                               Note: We must apply `VECTOR_ENFORCE()` on this column via
                               Rockset ingest transformation.

            """
            try:
                from rockset import RocksetClient
            except ImportError:
                raise ImportError(
                    "Could not import rockset client python package. "
                    "Please install it with `pip install rockset`."
                )

            if not isinstance(client, RocksetClient):
                raise ValueError(
                    f"client should be an instance of rockset.RocksetClient, "
                    f"got {type(client)}"
                )
            # TODO: check that `collection_name` exists in rockset. Create if not.
            self._client = client
            self._collection_name = collection_name
            self._embeddings = embeddings
            self._text_key = text_key
            self._embedding_key = embedding_key
            self._workspace = workspace

            try:
                self._client.set_application("langchain")
            except AttributeError:
                # ignore
                pass

        @property
        def embeddings(self) -> Embeddings:
            return self._embeddings

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

                    Args:
                texts: Iterable of strings to add to the vectorstore.
                metadatas: Optional list of metadatas associated with the texts.
                ids: Optional list of ids to associate with the texts.
                batch_size: Send documents in batches to rockset.

            Returns:
                List of ids from adding the texts into the vectorstore.

            """
            batch: list[dict] = []
            stored_ids = []

            for i, text in enumerate(texts):
                if len(batch) == batch_size:
                    stored_ids += self._write_documents_to_rockset(batch)
                    batch = []
                doc = {}
                if metadatas and len(metadatas) > i:
                    doc = deepcopy(metadatas[i])
                if ids and len(ids) > i:
                    doc["_id"] = ids[i]
                doc[self._text_key] = text
                doc[self._embedding_key] = self._embeddings.embed_query(text)
                batch.append(doc)
            if len(batch) > 0:
                stored_ids += self._write_documents_to_rockset(batch)
                batch = []
            return stored_ids

        @classmethod
        def from_texts(
            cls,
            texts: List[str],
            embedding: Embeddings,
            metadatas: Optional[List[dict]] = None,
            client: Any = None,
            collection_name: str = "",
            text_key: str = "",
            embedding_key: str = "",
            ids: Optional[List[str]] = None,
            batch_size: int = 32,
            **kwargs: Any,
        ) -> Rockset:
            """Create Rockset wrapper with existing texts.
            This is intended as a quicker way to get started.
            """

            # Sanitize inputs
            assert client is not None, "Rockset Client cannot be None"
            assert collection_name, "Collection name cannot be empty"
            assert text_key, "Text key name cannot be empty"
            assert embedding_key, "Embedding key cannot be empty"

            rockset = cls(client, embedding, collection_name, text_key, embedding_key)
            rockset.add_texts(texts, metadatas, ids, batch_size)
            return rockset

        # Rockset supports these vector distance functions.
        class DistanceFunction(Enum):
            COSINE_SIM = "COSINE_SIM"
            EUCLIDEAN_DIST = "EUCLIDEAN_DIST"
            DOT_PRODUCT = "DOT_PRODUCT"

            # how to sort results for "similarity"
            def order_by(self) -> str:
                if self.value == "EUCLIDEAN_DIST":
                    return "ASC"
                return "DESC"

        def similarity_search_with_relevance_scores(
            self,
            query: str,
            k: int = 4,
            distance_func: DistanceFunction = DistanceFunction.COSINE_SIM,
            where_str: Optional[str] = None,
            **kwargs: Any,
        ) -> List[Tuple[Document, float]]:
            """Perform a similarity search with Rockset

            Args:
                query (str): Text to look up documents similar to.
                distance_func (DistanceFunction): how to compute distance between two
                    vectors in Rockset.
                k (int, optional): Top K neighbors to retrieve. Defaults to 4.
                where_str (Optional[str], optional): Metadata filters supplied as a
                    SQL `where` condition string. Defaults to None.
                    eg. "price<=70.0 AND brand='Nintendo'"

                NOTE: Please do not let end-user to fill this and always be aware
                      of SQL injection.

            Returns:
                List[Tuple[Document, float]]: List of documents with their relevance score
            """
            return self.similarity_search_by_vector_with_relevance_scores(
                self._embeddings.embed_query(query),
                k,
                distance_func,
                where_str,
                **kwargs,
            )

        def similarity_search(
            self,
            query: str,
            k: int = 4,
            distance_func: DistanceFunction = DistanceFunction.COSINE_SIM,
            where_str: Optional[str] = None,
            **kwargs: Any,
        ) -> List[Document]:
            """Same as `similarity_search_with_relevance_scores` but
            doesn't return the scores.
            """
            return self.similarity_search_by_vector(
                self._embeddings.embed_query(query),
                k,
                distance_func,
                where_str,
                **kwargs,
            )

        def similarity_search_by_vector(
            self,
            embedding: List[float],
            k: int = 4,
            distance_func: DistanceFunction = DistanceFunction.COSINE_SIM,
            where_str: Optional[str] = None,
            **kwargs: Any,
        ) -> List[Document]:
            """Accepts a query_embedding (vector), and returns documents with
            similar embeddings."""

            docs_and_scores = self.similarity_search_by_vector_with_relevance_scores(
                embedding, k, distance_func, where_str, **kwargs
            )
            return [doc for doc, _ in docs_and_scores]

        def similarity_search_by_vector_with_relevance_scores(
            self,
            embedding: List[float],
            k: int = 4,
            distance_func: DistanceFunction = DistanceFunction.COSINE_SIM,
            where_str: Optional[str] = None,
            **kwargs: Any,
        ) -> List[Tuple[Document, float]]:
            """Accepts a query_embedding (vector), and returns documents with
            similar embeddings along with their relevance scores."""

            exclude_embeddings = True
            if "exclude_embeddings" in kwargs:
                exclude_embeddings = kwargs["exclude_embeddings"]
            q_str = self._build_query_sql(
                embedding, distance_func, k, where_str, exclude_embeddings
            )
            try:
                query_response = self._client.Queries.query(sql={"query": q_str})
            except Exception as e:
                logger.error("Exception when querying Rockset: %s\n", e)
                return []
            finalResult: list[Tuple[Document, float]] = []
            for document in query_response.results:
                metadata = {}
                assert isinstance(
                    document, dict
                ), "document should be of type `dict[str,Any]`. But found: `{}`".format(
                    type(document)
                )
                for k, v in document.items():
                    if k == self._text_key:
                        assert isinstance(v, str), (
                            "page content stored in column `{}` must be of type `str`. "
                            "But found: `{}`"
                        ).format(self._text_key, type(v))
                        page_content = v
                    elif k == "dist":
                        assert isinstance(v, float), (
                            "Computed distance between vectors must of type `float`. "
                            "But found {}"
                        ).format(type(v))
                        score = v
                    elif k not in ["_id", "_event_time", "_meta"]:
                        # These columns are populated by Rockset when documents are
                        # inserted. No need to return them in metadata dict.
                        metadata[k] = v
                finalResult.append(
                    (Document(page_content=page_content, metadata=metadata), score)
                )
            return finalResult

        def max_marginal_relevance_search(
            self,
            query: str,
            k: int = 4,
            fetch_k: int = 20,
            lambda_mult: float = 0.5,
            *,
            where_str: Optional[str] = None,
            **kwargs: Any,
        ) -> List[Document]:
            """Return docs selected using the maximal marginal relevance.

            Maximal marginal relevance optimizes for similarity to query AND diversity
            among selected documents.

            Args:
                query: Text to look up documents similar to.
                k: Number of Documents to return. Defaults to 4.
                fetch_k: Number of Documents to fetch to pass to MMR algorithm.
                distance_func (DistanceFunction): how to compute distance between two
                    vectors in Rockset.
                lambda_mult: Number between 0 and 1 that determines the degree
                            of diversity among the results with 0 corresponding
                            to maximum diversity and 1 to minimum diversity.
                            Defaults to 0.5.
                where_str: where clause for the sql query
            Returns:
                List of Documents selected by maximal marginal relevance.
            """
            query_embedding = self._embeddings.embed_query(query)
            initial_docs = self.similarity_search_by_vector(
                query_embedding,
                k=fetch_k,
                where_str=where_str,
                exclude_embeddings=False,
                **kwargs,
            )

            embeddings = [doc.metadata[self._embedding_key] for doc in initial_docs]

            selected_indices = maximal_marginal_relevance(
                np.array(query_embedding),
                embeddings,
                lambda_mult=lambda_mult,
                k=k,
            )

            # remove embeddings key before returning for cleanup to be consistent with
            #   other search functions
            for i in selected_indices:
                del initial_docs[i].metadata[self._embedding_key]

            return [initial_docs[i] for i in selected_indices]

        # Helper functions

        def _build_query_sql(
            self,
            query_embedding: List[float],
            distance_func: DistanceFunction,
            k: int = 4,
            where_str: Optional[str] = None,
            exclude_embeddings: bool = True,
        ) -> str:
            """Builds Rockset SQL query to query similar vectors to query_vector"""

            q_embedding_str = ",".join(map(str, query_embedding))
            distance_str = f"""{distance_func.value}({self._embedding_key}, \
    [{q_embedding_str}]) as dist"""
            where_str = f"WHERE {where_str}\n" if where_str else ""
            select_embedding = (
                f" EXCEPT({self._embedding_key})," if exclude_embeddings else ","
            )
            return f"""\
    SELECT *{select_embedding} {distance_str}
    FROM {self._workspace}.{self._collection_name}
    {where_str}\
    ORDER BY dist {distance_func.order_by()}
    LIMIT {str(k)}
    """

        def _write_documents_to_rockset(self, batch: List[dict]) -> List[str]:
            add_doc_res = self._client.Documents.add_documents(
                collection=self._collection_name, data=batch, workspace=self._workspace
            )
            return [doc_status._id for doc_status in add_doc_res.data]

        def delete_texts(self, ids: List[str]) -> None:
            """Delete a list of docs from the Rockset collection"""
            try:
                from rockset.models import DeleteDocumentsRequestData
            except ImportError:
                raise ImportError(
                    "Could not import rockset client python package. "
                    "Please install it with `pip install rockset`."
                )

            self._client.Documents.delete_documents(
                collection=self._collection_name,
                data=[DeleteDocumentsRequestData(id=i) for i in ids],
                workspace=self._workspace,
            )

        def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:
            try:
                if ids is None:
                    ids = []
                self.delete_texts(ids)
            except Exception as e:
                logger.error("Exception when deleting docs from Rockset: %s\n", e)
                return False

            return True

        async def adelete(
            self, ids: Optional[List[str]] = None, **kwargs: Any
        ) -> Optional[bool]:
            return await run_in_executor(None, self.delete, ids, **kwargs)