andreped/ReferenceBot

1from typing import Any2from typing import Iterable3from typing import List4from typing import Optional5 6from langchain.chat_models.fake import FakeListChatModel7from langchain.docstore.document import Document8from langchain.embeddings.base import Embeddings9from langchain.embeddings.fake import FakeEmbeddings as FakeEmbeddingsBase10from langchain.vectorstores import VectorStore11 12 13class FakeChatModel(FakeListChatModel):14    def __init__(self, **kwargs):15        responses = ["The answer is 42. SOURCES: 1, 2, 3, 4"]16        super().__init__(responses=responses, **kwargs)17 18 19class FakeEmbeddings(FakeEmbeddingsBase):20    def __init__(self, **kwargs):21        super().__init__(size=4, **kwargs)22 23 24class FakeVectorStore(VectorStore):25    """Fake vector store for testing purposes."""26 27    def __init__(self, texts: List[str]):28        self.texts: List[str] = texts29 30    def add_texts(self, texts: Iterable[str], metadatas: List[dict] | None = None, **kwargs: Any) -> List[str]:31        self.texts.extend(texts)32        return self.texts33 34    @classmethod35    def from_texts(36        cls,37        texts: List[str],38        embedding: Embeddings,39        metadatas: Optional[List[dict]] = None,40        **kwargs: Any,41    ) -> "FakeVectorStore":42        return cls(texts=list(texts))43 44    def similarity_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]:45        return [Document(page_content=text, metadata={"source": f"{i+1}-{1}"}) for i, text in enumerate(self.texts)]46