new-rag

New RAG Pipeline Command

Create a new RAG pipeline file with best practices baked in.

Workflow

1. Ask the user which vector store they want:

   • InMemory - Quick prototyping, no persistence
   • FAISS - Local, file-based persistence
   • Chroma - Local with server option
   • pgvector - PostgreSQL-based
   • Pinecone - Managed cloud service

2. Get filename from argument or ask user (default: rag_pipeline.py)

3. Generate the RAG file using the appropriate template below

4. Inform user about next steps (install dependencies, configure embeddings)

Templates

InMemory Template

"""
RAG Pipeline using InMemoryVectorStore

Install: pip install langchain langchain-openai
"""
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import OpenAIEmbeddings, ChatOpenAI


# 1. Load documents
def load_documents(urls: list[str]):
    """Load documents from URLs."""
    docs = []
    for url in urls:
        loader = WebBaseLoader(url)
        docs.extend(loader.load())
    return docs


# 2. Split into chunks
def split_documents(docs, chunk_size=1000, chunk_overlap=200):
    """Split documents into chunks for embedding."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        add_start_index=True,
    )
    return splitter.split_documents(docs)


# 3. Create vector store
def create_vectorstore(chunks):
    """Create vector store from document chunks."""
    embeddings = OpenAIEmbeddings()
    vectorstore = InMemoryVectorStore.from_documents(chunks, embeddings)
    return vectorstore


# 4. Build RAG chain
def build_rag_chain(vectorstore):
    """Build the RAG chain."""
    retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

    prompt = ChatPromptTemplate.from_template("""
Answer the question based only on the following context.
If you cannot answer based on the context, say "I don't know."

Context:
{context}

Question: {question}

Answer:
""")

    llm = ChatOpenAI(model="gpt-4o", temperature=0)

    def format_docs(docs):
        return "\n\n".join(doc.page_content for doc in docs)

    chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    return chain


def main():
    # Example usage
    urls = [
        "https://example.com/page1",
        "https://example.com/page2",
    ]

    print("Loading documents...")
    docs = load_documents(urls)
    print(f"Loaded {len(docs)} documents")

    print("Splitting documents...")
    chunks = split_documents(docs)
    print(f"Created {len(chunks)} chunks")

    print("Creating vector store...")
    vectorstore = create_vectorstore(chunks)

    print("Building RAG chain...")
    chain = build_rag_chain(vectorstore)

    # Query
    question = "What is this about?"
    print(f"\nQuestion: {question}")
    answer = chain.invoke(question)
    print(f"Answer: {answer}")


if __name__ == "__main__":
    main()