pip install streamlit langchain langchain-ollama langchain-community chromadb python-dotenv pypdf

pip install langchain-openai

OPENAI_API_KEY=your-api-key-here

ollama pull llama3.2
ollama pull qwen2.5

import streamlit as st
from dotenv import load_dotenv
from langchain.schema import HumanMessage, AIMessage, SystemMessage
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader
from langchain.chains import RetrievalQA
import os
import tempfile

load_dotenv()

def configure_page():
    st.set_page_config(
        page_title="PDF Chat with LangChain, Ollama, and Chroma",
        page_icon="🤖",
        layout="wide",
        initial_sidebar_state="expanded",
    )
    st.title("📄🤖 Chat with Your PDF using LangChain, Ollama, and Chroma")
    with st.expander("Check State"):
        st.write(st.session_state)

@st.cache_resource
def get_chat_model(model_name):
    if model_name == "gpt-3.5-turbo":
        from langchain_openai import ChatOpenAI
        return ChatOpenAI(
            api_key=os.getenv("OPENAI_API_KEY"),
            model=model_name,
            streaming=True,
        )
    return ChatOllama(model=model_name, streaming=True)

@st.cache_resource
def get_embeddings():
    return OllamaEmbeddings(model="mxbai-embed-large")

def load_pdf(uploaded_file):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
        tmp_file.write(uploaded_file.read())
        tmp_file_path = tmp_file.name
    loader = PyPDFLoader(tmp_file_path)
    documents = loader.load()
    os.unlink(tmp_file_path)
    return documents

def split_text(documents):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len
    )
    texts = text_splitter.split_documents(documents)
    return texts

def create_vector_store(texts, embeddings):
    chroma_persist_directory = os.path.join(tempfile.gettempdir(), "chroma_db")
    vector_store = Chroma.from_documents(
        documents=texts,
        embedding=embeddings,
        persist_directory=chroma_persist_directory
    )
    vector_store.persist()
    return vector_store

def display_chat_messages():
    for message in st.session_state.messages[1:]:
        if isinstance(message, HumanMessage):
            with st.chat_message("user"):
                st.write(message.content)
        elif isinstance(message, AIMessage):
            with st.chat_message("assistant"):
                st.write(message.content)

def handle_user_input(chat_model, retriever):
    if prompt := st.chat_input("Ask something about your PDF"):
        st.session_state.messages.append(HumanMessage(content=prompt))
        with st.chat_message("user"):
            st.write(prompt)

        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            try:
                response = chat_model.run(prompt)
                full_response = response
                message_placeholder.markdown(full_response)
                st.session_state.messages.append(AIMessage(content=full_response))
            except Exception as e:
                message_placeholder.markdown("❌ An error occurred while generating the response.")
                st.error(f"Error: {e}")

def main():
    configure_page()
    selected_model = handle_sidebar()
    chat_model = get_chat_model(selected_model)

    if "messages" not in st.session_state:
        st.session_state.messages = [
            SystemMessage(content="You are a helpful AI assistant.")
        ]

    if "vector_store" in st.session_state:
        retriever = st.session_state.vector_store.as_retriever()
        qa_chain = RetrievalQA.from_chain_type(
            llm=chat_model,
            chain_type="stuff",
            retriever=retriever,
            return_source_documents=False,
        )
        chat_model_with_retrieval = qa_chain
    else:
        chat_model_with_retrieval = chat_model

    display_chat_messages()

    retriever_instance = (
        st.session_state.vector_store.as_retriever()
        if "vector_store" in st.session_state else None
    )
    handle_user_input(chat_model_with_retrieval, retriever=retriever_instance)

if __name__ == "__main__":
    main()

streamlit run app.py