Toggle Menu
Prev Up Next

# Source code for langchain_community.vectorstores.inmemory

    import uuid
    from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

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

    from langchain_community.utils.math import cosine_similarity
    from langchain_community.vectorstores.utils import maximal_marginal_relevance

    class InMemoryVectorStore(VectorStore):
        """In-memory implementation of VectorStore using a dictionary.
        Uses numpy to compute cosine similarity for search.

        Args:
            embedding:  embedding function to use.
        """

        def __init__(self, embedding: Embeddings) -> None:
            self.store: Dict[str, Dict[str, Any]] = {}
            self.embedding = embedding

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

        def delete(self, ids: Optional[Sequence[str]] = None, **kwargs: Any) -> None:
            if ids:
                for _id in ids:
                    self.store.pop(_id, None)

        async def adelete(self, ids: Optional[Sequence[str]] = None, **kwargs: Any) -> None:
            self.delete(ids)

        def add_texts(
            self,
            texts: Iterable[str],
            metadatas: Optional[List[dict]] = None,
            **kwargs: Any,
        ) -> List[str]:
            ids = []
            vectors = self.embedding.embed_documents(list(texts))

            for i, text in enumerate(texts):
                doc_id = str(uuid.uuid4())
                ids.append(doc_id)
                self.store[doc_id] = {
                    "id": doc_id,
                    "vector": vectors[i],
                    "text": text,
                    "metadata": metadatas[i] if metadatas else {},
                }
            return ids

        async def aadd_texts(
            self,
            texts: Iterable[str],
            metadatas: Optional[List[dict]] = None,
            **kwargs: Any,
        ) -> List[str]:
            return self.add_texts(texts, metadatas, **kwargs)

        def similarity_search_with_score_by_vector(
            self,
            embedding: List[float],
            k: int = 4,
        ) -> List[Tuple[Document, float]]:
            docs_with_similarity = []
            for doc in self.store.values():
                similarity = float(cosine_similarity([embedding], [doc["vector"]]).item(0))
                docs_with_similarity.append(
                    (
                        Document(page_content=doc["text"], metadata=doc["metadata"]),
                        similarity,
                    )
                )
            docs_with_similarity.sort(key=lambda x: x[1], reverse=True)
            return docs_with_similarity[:k]

        def similarity_search_with_score(
            self,
            query: str,
            k: int = 4,
            **kwargs: Any,
        ) -> List[Tuple[Document, float]]:
            embedding = self.embedding.embed_query(query)
            docs = self.similarity_search_with_score_by_vector(
                embedding,
                k,
            )
            return docs

        async def asimilarity_search_with_score(
            self, query: str, k: int = 4, **kwargs: Any
        ) -> List[Tuple[Document, float]]:
            return self.similarity_search_with_score(query, k, **kwargs)

        def similarity_search_by_vector(
            self,
            embedding: List[float],
            k: int = 4,
            **kwargs: Any,
        ) -> List[Document]:
            docs_and_scores = self.similarity_search_with_score_by_vector(
                embedding,
                k,
            )
            return [doc for doc, _ in docs_and_scores]

        async def asimilarity_search_by_vector(
            self, embedding: List[float], k: int = 4, **kwargs: Any
        ) -> List[Document]:
            return self.similarity_search_by_vector(embedding, k, **kwargs)

        def similarity_search(
            self, query: str, k: int = 4, **kwargs: Any
        ) -> List[Document]:
            return [doc for doc, _ in self.similarity_search_with_score(query, k, **kwargs)]

        async def asimilarity_search(
            self, query: str, k: int = 4, **kwargs: Any
        ) -> List[Document]:
            return self.similarity_search(query, k, **kwargs)

        def max_marginal_relevance_search_by_vector(
            self,
            embedding: List[float],
            k: int = 4,
            fetch_k: int = 20,
            lambda_mult: float = 0.5,
            **kwargs: Any,
        ) -> List[Document]:
            docs_with_similarity = []
            for doc in self.store.values():
                similarity = float(cosine_similarity([embedding], [doc["vector"]]).item(0))
                docs_with_similarity.append(
                    (
                        doc,
                        similarity,
                    )
                )
            docs_with_similarity.sort(key=lambda x: x[1], reverse=True)
            prefetch_hits = docs_with_similarity[:fetch_k]

            mmr_chosen_indices = maximal_marginal_relevance(
                np.array(embedding, dtype=np.float32),
                [doc["vector"] for doc, _ in prefetch_hits],
                k=k,
                lambda_mult=lambda_mult,
            )
            return [
                Document(
                    page_content=prefetch_hits[idx][0]["text"],
                    metadata=prefetch_hits[idx][0]["metadata"],
                )
                for idx in mmr_chosen_indices
            ]

        def max_marginal_relevance_search(
            self,
            query: str,
            k: int = 4,
            fetch_k: int = 20,
            lambda_mult: float = 0.5,
            **kwargs: Any,
        ) -> List[Document]:
            embedding_vector = self.embedding.embed_query(query)
            return self.max_marginal_relevance_search_by_vector(
                embedding_vector,
                k,
                fetch_k,
                lambda_mult=lambda_mult,
            )

        @classmethod
        def from_texts(
            cls,
            texts: List[str],
            embedding: Embeddings,
            metadatas: Optional[List[dict]] = None,
            **kwargs: Any,
        ) -> "InMemoryVectorStore":
            store = cls(
                embedding=embedding,
            )
            store.add_texts(texts=texts, metadatas=metadatas)
            return store

        @classmethod
        async def afrom_texts(
            cls,
            texts: List[str],
            embedding: Embeddings,
            metadatas: Optional[List[dict]] = None,
            **kwargs: Any,
        ) -> "InMemoryVectorStore":
            return cls.from_texts(texts, embedding, metadatas, **kwargs)