andreped/ReferenceBot

1from typing import List2from typing import Type3 4from langchain.docstore.document import Document5from langchain.embeddings import OpenAIEmbeddings6from langchain.embeddings.base import Embeddings7from langchain.vectorstores import VectorStore8from langchain.vectorstores.faiss import FAISS9 10from knowledge_gpt.core.debug import FakeEmbeddings11from knowledge_gpt.core.debug import FakeVectorStore12from knowledge_gpt.core.parsing import File13 14 15class FolderIndex:16    """Index for a collection of files (a folder)"""17 18    def __init__(self, files: List[File], index: VectorStore):19        self.name: str = "default"20        self.files = files21        self.index: VectorStore = index22 23    @staticmethod24    def _combine_files(files: List[File]) -> List[Document]:25        """Combines all the documents in a list of files into a single list."""26 27        all_texts = []28        for file in files:29            for doc in file.docs:30                doc.metadata["file_name"] = file.name31                doc.metadata["file_id"] = file.id32                all_texts.append(doc)33 34        return all_texts35 36    @classmethod37    def from_files(cls, files: List[File], embeddings: Embeddings, vector_store: Type[VectorStore]) -> "FolderIndex":38        """Creates an index from files."""39 40        all_docs = cls._combine_files(files)41 42        index = vector_store.from_documents(43            documents=all_docs,44            embedding=embeddings,45        )46 47        return cls(files=files, index=index)48 49 50def embed_files(files: List[File], embedding: str, vector_store: str, **kwargs) -> FolderIndex:51    """Embeds a collection of files and stores them in a FolderIndex."""52 53    supported_embeddings: dict[str, Type[Embeddings]] = {54        "openai": OpenAIEmbeddings,55        "debug": FakeEmbeddings,56    }57    supported_vector_stores: dict[str, Type[VectorStore]] = {58        "faiss": FAISS,59        "debug": FakeVectorStore,60    }61 62    if embedding in supported_embeddings:63        _embeddings = supported_embeddings[embedding](**kwargs)64    else:65        raise NotImplementedError(f"Embedding {embedding} not supported.")66 67    if vector_store in supported_vector_stores:68        _vector_store = supported_vector_stores[vector_store]69    else:70        raise NotImplementedError(f"Vector store {vector_store} not supported.")71 72    return FolderIndex.from_files(files=files, embeddings=_embeddings, vector_store=_vector_store)73