app.py

import os
from typing import List
import uuid
import chainlit as cl
from chainlit.types import AskFileResponse
from langchain.memory import ConversationBufferMemory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.document_loaders import PyMuPDFLoader, TextLoader
from langchain.prompts import MessagesPlaceholder
from langchain.prompts import ChatPromptTemplate
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_experimental.text_splitter import SemanticChunker
from langchain_qdrant import QdrantVectorStore
from langchain_core.documents import Document
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from langchain_openai import ChatOpenAI
from embedding_model import get_embeddings_openai_text_3_large,get_embeddings_snowflake_arctic_embed_l
from pdfloader import PDFLoaderWrapper
from langchain_core.runnables.history import RunnableWithMessageHistory
from chainlit.input_widget import Select, Switch, Slider
from dotenv import load_dotenv
from langchain_huggingface import HuggingFaceEmbeddings

load_dotenv()

BOR_FILE_PATH = "https://www.whitehouse.gov/wp-content/uploads/2022/10/Blueprint-for-an-AI-Bill-of-Rights.pdf"
NIST_FILE_PATH = "https://nvlpubs.nist.gov/nistpubs/ai/NIST.AI.600-1.pdf"
SMALL_DOC = "https://arxiv.org/pdf/1908.10084"  # 11 pages Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks
documents_to_preload = [
    BOR_FILE_PATH,
    NIST_FILE_PATH
    # SMALL_DOC
]
collection_name = "ai-safety"

welcome_message = """
Welcome to the chatbot to clarify all your AI Safety related queries.:
Now preloading below documents:
    1. Blueprint for an AI Bill of Rights
    2. NIST AI Standards
Please wait for a moment to load the documents.
"""
chat_model_name = "gpt-4o"
embedding_model_name = "Snowflake/snowflake-arctic-embed-l"
chat_model = ChatOpenAI(model=chat_model_name, temperature=0)

async def connect_to_qdrant():
    embedding_model = HuggingFaceEmbeddings(model_name=embedding_model_name)
    qdrant_url = os.environ["QDRANT_URL"]
    qdrant_api_key = os.environ["QDRANT_API_KEY"]
    collection_name = os.environ["COLLECTION_NAME"]
    qdrant_client = QdrantClient(url=qdrant_url,api_key=qdrant_api_key)
    vector_store = QdrantVectorStore(
        client=qdrant_client,
        collection_name=collection_name,
        embedding=embedding_model,
    )
    return vector_store.as_retriever()

def initialize_vectorstore(
    collection_name: str,
    embedding_model,
    dimension,
    distance_metric: Distance = Distance.COSINE,
):
    client = QdrantClient(":memory:")
    client.create_collection(
        collection_name=collection_name,
        vectors_config=VectorParams(size=dimension, distance=distance_metric),
    )

    vector_store = QdrantVectorStore(
        client=client,
        collection_name=collection_name,
        embedding=embedding_model,
    )
    return vector_store

def get_text_splitter(strategy, embedding_model):
    if strategy == "semantic":
        return SemanticChunker(
            embedding_model,
            breakpoint_threshold_type="percentile",
            breakpoint_threshold_amount=90,
        )

def process_file(file: AskFileResponse, text_splitter):
    if file.type == "text/plain":
        Loader = TextLoader
    elif file.type == "application/pdf":
        Loader = PyMuPDFLoader

    loader = Loader(file.path)
    documents = loader.load()
    title = documents[0].metadata.get("title")
    docs = text_splitter.split_documents(documents)
    for i, doc in enumerate(docs):
        doc.metadata["source"] = f"source_{i}"
        doc.metadata["title"] = title
    return docs

def populate_vectorstore(vector_store, docs: List[Document]):
    vector_store.add_documents(docs)
    return vector_store

def create_history_aware_retriever_self(chat_model, retriever):
    contextualize_q_system_prompt = (
        "Given a chat history and the latest user question which might reference context in the chat history, "
        "formulate a standalone question which can be understood without the chat history. Do NOT answer the question, "
        "just reformulate it if needed and otherwise return it as is."
    )
    contextualize_q_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", contextualize_q_system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )
    return create_history_aware_retriever(chat_model, retriever, contextualize_q_prompt)

def create_qa_chain(chat_model):
    qa_system_prompt = (
        "You are an helpful assistant named 'Shield' and your task is to answer any questions related to AI Safety for the given context."
        "Use the following pieces of retrieved context to answer the question."
        # "If any questions asked outside AI Safety context, just say that you are a specialist in AI Safety and can't answer that."
        # f"When introducing you, just say that you are an AI assistant powered by embedding model {embedding_model_name} and chat model {chat_model_name} and your knowledge is limited to 'Blueprint for an AI Bill of Rights' and 'NIST AI Standards' documents."
        "If you don't know the answer, just say that you don't know.\n\n"
        "{context}"
    )
    qa_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", qa_system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )
    return create_stuff_documents_chain(chat_model, qa_prompt)


def create_rag_chain(chat_model, retriever):
    history_aware_retriever = create_history_aware_retriever_self(chat_model, retriever)
    question_answer_chain = create_qa_chain(chat_model)
    return create_retrieval_chain(history_aware_retriever, question_answer_chain)


def create_session_id():
    session_id = str(uuid.uuid4())
    return session_id


@cl.on_chat_start
async def start():
    # cl.user_session.set("memory", conversation_buffer_memory)
    msg = cl.Message(content=welcome_message)
    await msg.send()

    # Create a session id
    session_id = create_session_id()
    cl.user_session.set("session_id", session_id)

    # Preserve chat history
    conversation_buffer_memory = ConversationBufferMemory(
        memory_key="chat_history",
        output_key="answer",
        chat_memory=ChatMessageHistory(),
        return_messages=True,
    )

    # todo: if logged in user is admin then allow them to upload new pdfs.

    # # Embedding model
    # # embedding_model, dimension = get_embeddings_openai_text_3_large()
    # embedding_model, dimension = get_embeddings_snowflake_arctic_embed_l()
    # msg.content = "Embedding model loaded"
    # await msg.update()
    # cl.user_session.set("embedding_model", embedding_model)
    # cl.user_session.set("dimension", dimension)

    # # Pdf loader
    # pdf_loader = PDFLoaderWrapper(
    #     documents_to_preload, PDFLoaderWrapper.LoaderType.PYMUPDF
    # )
    # msg.content = "Embedding model loaded"
    # await msg.update()
    # cl.user_session.set("pdf_loader", pdf_loader)
    # documents = await pdf_loader.aload()

    # text_splitter = get_text_splitter("semantic", embedding_model)

    # chunked_docs = text_splitter.split_documents(documents)

    # vector_store = initialize_vectorstore(
    #     collection_name, embedding_model, dimension=dimension
    # )

    # vector_store = populate_vectorstore(vector_store, chunked_docs)

    retriever = await connect_to_qdrant()

    rag_chain = create_rag_chain(chat_model, retriever)

    store = {}

    def get_session_history(session_id: str) -> BaseChatMessageHistory:
        if session_id not in store:
            store[session_id] = ChatMessageHistory()
        return store[session_id]

    conversational_rag_chain = RunnableWithMessageHistory(
        rag_chain,
        get_session_history,
        input_messages_key="input",
        history_messages_key="chat_history",
        output_messages_key="answer",
    )

    # Let the user know that the system is ready
    msg.content = msg.content + "\nReady to answer your questions!"
    await msg.update()

    cl.user_session.set("conversational_rag_chain", conversational_rag_chain)


@cl.on_message
async def main(message: cl.Message):
    session_id = cl.user_session.get("session_id")
    conversational_rag_chain = cl.user_session.get("conversational_rag_chain")

    response =  await conversational_rag_chain.ainvoke(
        {"input": message.content},
        config={"configurable": {"session_id": session_id},
                "callbacks":[cl.AsyncLangchainCallbackHandler()]},         
    )
    answer = response["answer"]

    source_documents = response["context"]
    text_elements = []
    unique_pages = set()

    if source_documents:

        for source_idx, source_doc in enumerate(source_documents):
            source_name = f"source_{source_idx+1}"
            page_number = source_doc.metadata['page']
            #page_number = source_doc.metadata.get('page', "NA")  # NA or any default value
            page = f"Page {page_number}"
            text_element_content = source_doc.page_content
            text_element_content = text_element_content if text_element_content != "" else "No Content"
            #text_elements.append(cl.Text(content=text_element_content, name=source_name))
            if page not in unique_pages:
                unique_pages.add(page)
                text_elements.append(cl.Text(content=text_element_content, name=page))
            #text_elements.append(cl.Text(content=text_element_content, name=page))
        source_names = [text_el.name for text_el in text_elements]
        
        if source_names:
            answer += f"\n\n Sources:{', '.join(source_names)}"
        else:
            answer += "\n\n No sources found"

    await cl.Message(content=answer, elements=text_elements).send()

if __name__ == "__main__":
    from chainlit.cli import run_chainlit

    run_chainlit(__file__)