Basic RAG Patterns

import streamlit as st
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from langchain_chroma import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters.sentence_transformers import SentenceTransformersTokenTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# Initialize embedding model
embedding_model = GoogleGenerativeAIEmbeddings(
    model="models/embedding-001"
)

# Initialize vector database
db = Chroma(
    collection_name="pharma_database",
    embedding_function=embedding_model,
    persist_directory='./pharma_db'
)

def format_docs(docs):
    """Format document list into single string."""
    return "\n\n".join(doc.page_content for doc in docs)

def add_documents_to_db(uploaded_files):
    """Process and add PDF files to vector database."""
    for uploaded_file in uploaded_files:
        # Save uploaded file temporarily
        temp_path = f"./temp/{uploaded_file.name}"
        with open(temp_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        
        # Load PDF
        loader = PyPDFLoader(temp_path)
        documents = loader.load()
        
        # Extract metadata and content
        doc_metadata = [doc.metadata for doc in documents]
        doc_content = [doc.page_content for doc in documents]
        
        # Split into chunks
        text_splitter = SentenceTransformersTokenTextSplitter(
            model_name="sentence-transformers/all-mpnet-base-v2",
            chunk_size=100,
            chunk_overlap=50
        )
        chunks = text_splitter.create_documents(doc_content, doc_metadata)
        
        # Add to database
        db.add_documents(chunks)

def run_rag_chain(query):
    """Execute RAG chain for given query."""
    # Create retriever with similarity search
    retriever = db.as_retriever(
        search_type="similarity",
        search_kwargs={'k': 5}
    )
    
    # Define prompt template
    prompt_template = ChatPromptTemplate.from_template("""
    You are a knowledgeable assistant.
    Answer the question based only on the following context:
    {context}

    Question: {question}

    Provide a detailed and accurate answer based on the context.
    Don't make up information not present in the context.
    """)
    
    # Initialize LLM
    llm = ChatGoogleGenerativeAI(
        model="gemini-1.5-pro",
        temperature=0.7
    )
    
    # Create RAG chain
    rag_chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt_template
        | llm
        | StrOutputParser()
    )
    
    # Execute chain
    return rag_chain.invoke(query)

# Streamlit UI
st.title("📚 Basic RAG System")

# Query input
query = st.text_area(
    "Enter your question:",
    placeholder="What would you like to know?"
)

if st.button("Submit"):
    if query:
        with st.spinner("Searching and generating answer...":
            result = run_rag_chain(query)
            st.write(result)
    else:
        st.warning("Please enter a question")

# Sidebar for document upload
with st.sidebar:
    st.header("Document Upload")
    uploaded_files = st.file_uploader(
        "Upload PDF documents",
        type=["pdf"],
        accept_multiple_files=True
    )
    
    if st.button("Process Documents"):
        if uploaded_files:
            with st.spinner("Processing documents..."):
                add_documents_to_db(uploaded_files)
                st.success("Documents added successfully!")
        else:
            st.warning("Please upload files first")

prompt = ChatPromptTemplate.from_template("""
Answer the question based on the following context:

{context}

Question: {question}

Answer:
""")