Captainspa/grounded-code

1# This file is for running a retrieval augmented generation on an existing vector db2from langchain.memory import ConversationBufferMemory3from langchain_core.prompts import ChatPromptTemplate4from langchain_core.messages import SystemMessage5from langchain_core.runnables import RunnableLambda, RunnablePassthrough6from langchain_core.output_parsers import StrOutputParser7from langchain_core.documents import Document8from operator import itemgetter9from helpers import clean_docs10 11# These are the urls that get ingested as documents. Should be a list of strings.12DEFAULT_URLS = [13    "https://www.gradio.app/guides/blocks-and-event-listeners",14]15# This is what you want to name the collection of documents (it will create a folder in the repo with this name)16DEFAULT_COLLECTION_NAME = "gradio_blocks_listeners_collection"17# This is the question you want to ask (retriever will choose chunks of the documents as context to answer the question)18DEFAULT_QUESTION = "How would I implement gradio components that would allow a button that says 'Add sources' and allow the user to paste urls into a text box so they get saved as sources?"19 20def format_docs(docs: list[Document]) -> str:21    """22    Formats the list of documents into a single string.23    Used to format the docs into a string for context that is passed to the LLM.24    """25    return "\n\n".join(doc.page_content for doc in docs)26 27def get_rag_template():28    """29    Fetches the RAG template for the prompt.30    This template expects to be passed values for both context and question.31    """32    template = """Answer the question based only on the following context:33    {context}34 35    Question: {question}36    """37    rag_prompt_template = ChatPromptTemplate.from_template(template)38    rag_prompt_template.messages.insert(0, 39        SystemMessage(40            content="You are an AI programming assistant. Use the document excerpts to respond to the best of your ability."41        )42    )43    return rag_prompt_template44 45def get_chain(retriever, llm, memory: ConversationBufferMemory | None = None):46    """47    Input: retriever (contains vectorstore with documents) and llm48    Returns a chain for the RAG pipeline.49    Can be invoked with a question, like `chain.invoke("How do I do x task using this framework?")` to get a response.50    """51    # Get prompt template52    rag_prompt_template = get_rag_template()53    # Set memory54    if memory is None:55        memory = ConversationBufferMemory(56        return_messages=True, input_key="question", output_key="answer"57    )58    # Load memory59    # This adds a "memory" key to the input object60    loaded_memory = RunnablePassthrough.assign(61        chat_history=RunnableLambda(memory.load_memory_variables) | itemgetter("history"),62    )63    chain = (64        {"context": retriever | format_docs, "question": RunnablePassthrough()}65        | loaded_memory66        | rag_prompt_template67        | llm68        | StrOutputParser()69    )70    return chain71 72def main(url: str | list[str] = DEFAULT_URLS, collection_name: str = DEFAULT_COLLECTION_NAME, question: str = DEFAULT_QUESTION):73    import embeddings74    # from models import get_openai_embedder_small, get_claude_sonnet75    from models import get_openai_embedder_large, get_claude_opus76    embedder = get_openai_embedder_large()77    llm = get_claude_opus()78    if isinstance(url, str):79        url = [url]80    vectorstore = None81    for i in range(len(url)):82        docs = embeddings.documents_from_url(url[i])83        docs = clean_docs(docs)84        chunked_docs = embeddings.split_documents(docs)85        if i == 0:86            vectorstore = embeddings.create_vectorstore(chunked_docs, embedder, collection_name)87        else:88            assert vectorstore is not None, "Vectorstore not initialized"89            vectorstore.add_documents(chunked_docs)90    retriever = vectorstore.as_retriever()91    # Can add optional arguments like search_kwargs={"score_threshold": 0.5}92    chain = get_chain(retriever, llm)93    output = chain.invoke(question)94    return output95 96if __name__ == "__main__":97    main()