SUTRA Document RAG Chatbot - Streamlit App with LangChain

import streamlit as st
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain.schema import HumanMessage
from langchain.callbacks.base import BaseCallbackHandler
import PyPDF2
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
api_key = os.getenv("SUTRA_API_KEY")

# Page configuration
st.set_page_config(
    page_title="Sutra Document RAG Chatbot",
    page_icon="📚",
    layout="wide"
)

# Streaming callback handler
class StreamHandler(BaseCallbackHandler):
    def __init__(self, container, initial_text=""):
        self.container = container
        self.text = initial_text

    def on_llm_new_token(self, token: str, **kwargs):
        self.text += token
        self.container.markdown(self.text)

# Initialize the ChatOpenAI model - base instance for caching
@st.cache_resource
def get_base_chat_model():
    return ChatOpenAI(
        api_key=os.getenv("SUTRA_API_KEY"),
        base_url="https://api.two.ai/v2",
        model="sutra-v2",
        temperature=0.7,
    )

# Create a streaming version of the model with callback handler
def get_streaming_chat_model(callback_handler=None):
    return ChatOpenAI(
        api_key=os.getenv("SUTRA_API_KEY"),
        base_url="https://api.two.ai/v2",
        model="sutra-v2",
        temperature=0.7,
        streaming=True,
        callbacks=[callback_handler] if callback_handler else None
    )

# Sidebar for document upload
st.sidebar.image("https://framerusercontent.com/images/3Ca34Pogzn9I3a7uTsNSlfs9Bdk.png", use_container_width=True)
with st.sidebar:
    st.title("📚 Sutra RAG Chatbot")
    uploaded_file = st.file_uploader("Upload a PDF document", type="pdf")
    st.divider()
    st.markdown("### About Sutra RAG Chatbot")
    st.markdown("This app uses Retrieval-Augmented Generation to answer questions based on uploaded PDF documents using the Sutra LLM.")

# Main app title
st.markdown(
    f'<h1><img src="https://framerusercontent.com/images/9vH8BcjXKRcC5OrSfkohhSyDgX0.png" width="60"/> Sutra Document RAG Chatbot 🤖</h1>',
    unsafe_allow_html=True
)

# Initialize session state for messages and vector store
if "messages" not in st.session_state:
    st.session_state.messages = []
if "vector_store" not in st.session_state:
    st.session_state.vector_store = None

# Process uploaded PDF
if uploaded_file is not None:
    try:
        # Read and extract text from PDF
        pdf_reader = PyPDF2.PdfReader(uploaded_file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text() or ""

        # Split text into chunks
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        documents = text_splitter.split_text(text)

        # Create embeddings and vector store
        embeddings = OpenAIEmbeddings(api_key=os.getenv("SUTRA_API_KEY"), base_url="https://api.two.ai/v2")
        st.session_state.vector_store = FAISS.from_texts(documents, embeddings)

        st.success("Document processed successfully!")
    except Exception as e:
        st.error(f"Error processing PDF: {str(e)}")

# Display chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# Chat input
user_input = st.chat_input("Ask a question about the document...")

# Process user input
if user_input and st.session_state.vector_store is not None:
    # Add user message to chat
    st.session_state.messages.append({"role": "user", "content": user_input})

    # Display user message
    with st.chat_message("user"):
        st.write(user_input)

    # Generate response
    try:
        # Create message placeholder
        with st.chat_message("assistant"):
            response_placeholder = st.empty()

            # Create a stream handler
            stream_handler = StreamHandler(response_placeholder)

            # Get streaming model with handler
            chat = get_streaming_chat_model(stream_handler)

            # Set up RetrievalQA chain
            qa_chain = RetrievalQA.from_chain_type(
                llm=chat,
                chain_type="stuff",
                retriever=st.session_state.vector_store.as_retriever()
            )

            # Generate response using RAG
            response = qa_chain.invoke({"query": user_input})
            answer = response["result"]

            # Add assistant response to chat history
            st.session_state.messages.append({"role": "assistant", "content": answer})

    except Exception as e:
        st.error(f"Error: {str(e)}")
        if "API key" in str(e):
            st.error("Please check your Sutra API key in the environment variables.")
elif user_input and st.session_state.vector_store is None:
    st.error("Please upload a PDF document first.")