app.py

import os
import streamlit as st
from dotenv import load_dotenv
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import llamacpp
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler
from langchain.chains import create_history_aware_retriever, create_retrieval_chain, ConversationalRetrievalChain
from langchain.document_loaders import TextLoader
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.chat_message_histories.streamlit import StreamlitChatMessageHistory
from langchain.prompts import PromptTemplate
from langchain.vectorstores import Chroma
from utills import load_txt_documents, split_docs, load_uploaded_documents, retriever_from_chroma
from langchain.text_splitter import TokenTextSplitter, RecursiveCharacterTextSplitter
from langchain_community.document_loaders.directory import DirectoryLoader

script_dir = os.path.dirname(os.path.abspath(__file__))
data_path = os.path.join(script_dir, "data/")
model_path = os.path.join(script_dir, 'qwen2-0_5b-instruct-q4_0.gguf')
store = {}

model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}

hf = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)

def get_vectorstore(text_chunks):
    model_name = "sentence-transformers/all-mpnet-base-v2"
    model_kwargs = {'device': 'cpu'}
    encode_kwargs = {'normalize_embeddings': True}
    hf = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )

    vectorstore = Chroma.from_documents(documents=text_chunks, embedding=hf, persist_directory="docs/chroma/")
    return vectorstore

def get_pdf_text(pdf_docs):
    document_loader = DirectoryLoader(pdf_docs)
    return document_loader.load()

def get_text_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        separator="\n",
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len
    )
    chunks = text_splitter.split_text(text)
    return chunks

def create_conversational_rag_chain(vectorstore):
    
    script_dir = os.path.dirname(os.path.abspath(__file__))
    model_path = os.path.join(script_dir, 'qwen2-0_5b-instruct-q4_0.gguf')
    
    retriever = vectorstore.as_retriever(search_type='mmr', search_kwargs={"k": 7})

    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])

    llm = llamacpp.LlamaCpp(
        model_path=os.path.join(model_path),
        n_gpu_layers=1,
        temperature=0.1,
        top_p=0.9,
        n_ctx=22000,
        max_tokens=200,
        repeat_penalty=1.7,
        callback_manager=callback_manager,
        verbose=False,
    )

    contextualize_q_system_prompt = """Given a context, chat history and the latest user question
    which maybe 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."""

    ha_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_system_prompt)

    qa_system_prompt = """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as informative as possible, be polite and formal.\n{context}"""

    qa_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", qa_system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )

    question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)

    rag_chain = create_retrieval_chain(ha_retriever, question_answer_chain)
    msgs = StreamlitChatMessageHistory(key="special_app_key")

    conversation_chain = RunnableWithMessageHistory(
        rag_chain,
        lambda session_id: msgs,
        input_messages_key="input",
        history_messages_key="chat_history",
        output_messages_key="answer",
    )
    return conversation_chain

def main():
    """Main function for the Streamlit app."""
    # Initialize chat history if not already present in session state

    documents = []
    
    script_dir = os.path.dirname(os.path.abspath(__file__))
    data_path = os.path.join(script_dir, "data/")
    for filename in os.listdir(data_path):

    if filename.endswith('.txt'):

        file_path = os.path.join(data_path, filename)

        documents = TextLoader(file_path).load()

        documents.extend(documents)
    
    
    docs = split_docs(documents, 350, 40)

    
    
    
    
    

    vectorstore = get_vectorstore(docs)
    
    msgs = st.session_state.get("chat_history", StreamlitChatMessageHistory(key="special_app_key"))
    chain_with_history = create_conversational_rag_chain(vectorstore)

    st.title("Conversational RAG Chatbot")

    if prompt := st.chat_input():
        st.chat_message("human").write(prompt)

        # Prepare the input dictionary with the correct keys
        input_dict = {"input": prompt, "chat_history": msgs.messages}
        config = {"configurable": {"session_id": "any"}}

        # Process user input and handle response
        response = chain_with_history.invoke(input_dict, config)
        st.chat_message("ai").write(response["answer"])

        # Display retrieved documents (if any and present in response)
        if "docs" in response and response["documents"]:
            for index, doc in enumerate(response["documents"]):
                with st.expander(f"Document {index + 1}"):
                    st.write(doc)

    # Update chat history in session state
    st.session_state["chat_history"] = msgs

if __name__ == "__main__":
    main()