test_similarity.py — langchain Source File

Source Code

from collections.abc import Iterable
from typing import Any

from typing_extensions import override

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings, FakeEmbeddings
from langchain_core.example_selectors import (
    MaxMarginalRelevanceExampleSelector,
    SemanticSimilarityExampleSelector,
)
from langchain_core.vectorstores import VectorStore


class DummyVectorStore(VectorStore):
    def __init__(self, init_arg: str | None = None):
        self.texts: list[str] = []
        self.metadatas: list[dict] = []
        self._embeddings: Embeddings | None = None
        self.init_arg = init_arg

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

    @override
    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: list[dict] | None = None,
        **kwargs: Any,
    ) -> list[str]:
        self.texts.extend(texts)
        if metadatas:
            self.metadatas.extend(metadatas)
        return ["dummy_id"]

    @override
    def similarity_search(
        self, query: str, k: int = 4, **kwargs: Any
    ) -> list[Document]:
        return [
            Document(
                page_content=query, metadata={"query": query, "k": k, "other": "other"}
            )
        ] * k

    @override
    def max_marginal_relevance_search(
        self,
        query: str,
        k: int = 4,
        fetch_k: int = 20,
        lambda_mult: float = 0.5,
        **kwargs: Any,
    ) -> list[Document]:
        return [
            Document(
                page_content=query,
                metadata={"query": query, "k": k, "fetch_k": fetch_k, "other": "other"},
// ... (188 more lines)