Captainspa/grounded-code

1# This file is for testing the use of the embedding and creates a local vectorstore2 3from langchain.text_splitter import RecursiveCharacterTextSplitter4from langchain_community.document_loaders import WebBaseLoader, TextLoader, NotebookLoader5from langchain_core.documents import Document6from langchain_community.vectorstores import Chroma7from chromadb.config import Settings8from os.path import exists9 10def reformat_url(url):11    if url.count('/blob/') != 1:12        raise ValueError("Use the raw link to the notebook file")13    # remove /blob from the url14    url = url.replace('/blob', '')15    # append raw to the url16    url = url.replace('github.com', 'raw.githubusercontent.com')17    return url18 19def load_local_notebook(url: str):20    assert url[-6:] == ".ipynb", "Make sure the link is a valid notebook"21    import requests22    # URL to the raw text of the .ipynb file23    if not url.startswith("https://raw.github"):24        url = reformat_url(url)25    assert url.startswith("https://raw.githubuser")26    # Download the notebook as a file27    response = requests.get(url)28    local_file = "temp.ipynb"29    with open(local_file, 'wb') as f:30        f.write(response.content)31    loader = NotebookLoader(32        local_file,33        include_outputs=True,34        max_output_length=20,35        remove_newline=True,36    )37    return loader38 39def documents_from_url(url: str) -> list[Document]:40    """41    Load documents from a URL, return List[Document]42    """43    assert True, "Make sure the link is valid"44    print("Indexing url:", url)45    if url[-6:] == ".ipynb":46        loader = load_local_notebook(url)47    else:48        loader = WebBaseLoader(url)49    docs = loader.load()50    if not docs:51        raise ValueError(f"No documents found at {url}")52    return docs53 54def documents_from_text_file(filepath: str = "sample.txt") -> list[Document]:55    """56    Load documents from a string, return List[Document]57    """58    loader = TextLoader(filepath)59    docs = loader.load()60    return docs61 62def split_documents(docs: list[Document], chunk_size=4000, chunk_overlap=200) -> list[Document]:63    """64    Split documents into chunks, return List[Document]65    """66    chunked_docs = RecursiveCharacterTextSplitter(67        chunk_size=chunk_size, chunk_overlap=chunk_overlap68    ).split_documents(docs)69    return chunked_docs70 71def create_vectorstore(docs: list[Document], embedder, collection_name = "test_collection") -> None:72    """73    Create a vectorstore from documents74    """75    is_local = False76    if exists(collection_name):77        print("Note: Collection seems to already exist! Not adding documents to the collection.")78        is_local = True79    vectorstore = Chroma(80        collection_name=collection_name, 81        embedding_function=embedder,82        persist_directory=collection_name,83        client_settings= Settings(anonymized_telemetry=False, is_persistent=True),84    )85    if not is_local:86        vectorstore.add_documents(docs)87    return vectorstore88 89def main():90    from models import get_openai_embedder91    embedder = get_openai_embedder()92    url = "https://python.langchain.com/docs/expression_language/cookbook/retrieval"93    docs = documents_from_url(url)94    chunked_docs = split_documents(docs)95    vectorstore = create_vectorstore(chunked_docs, embedder)96    output = vectorstore.similarity_search("How would I use memory in a function?", k=1)97    print(output)98    return output