app.py

import streamlit as st
from streaming import StreamHandler
import utils
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler 
from retriever import retriever
from langchain.chains import RetrievalQA
from langchain.llms import LlamaCpp
from dotenv import load_dotenv

class CustomDataChatbot:
    def __init__(self):
        # Initialize session state variables, including messages
        st.session_state.messages = []

    @st.spinner('Analyzing documents..')
    def setup_qa_chain(self):
        # Setup memory for contextual conversation        
        # memory = ConversationBufferMemory(
        #     memory_key='chat_history',
        #     return_messages=True
        # )

        callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
        # Setup LLM and QA chain
        llm = LlamaCpp(model_path="./models/openhermes-2.5-neural-chat-7b-v3-1-7b.Q5_K_M.gguf",
            temperature=0.34,
            max_tokens=4000,
            n_ctx=4096,
            top_p=1,
            callback_manager=callback_manager,
            verbose=True)

        # qa_chain = ConversationalRetrievalChain.from_llm(llm, retriever=retriever(), memory=memory, verbose=True)


        return RetrievalQA.from_chain_type( llm, retriever= retriever())

    @utils.enable_chat_history
    def main(self):
        load_dotenv()
        st.set_page_config(page_title="ChatPDF", page_icon="📄")
        st.header('Chat with your documents')
        st.write('Has access to custom documents and can respond to user queries by referring to the content within those documents')
        st.write('[![view source code ](https://img.shields.io/badge/view_source_code-gray?logo=github)](https://github.com/shashankdeshpande/langchain-chatbot/blob/master/pages/4_%F0%9F%93%84_chat_with_your_documents.py)')

        user_query = st.chat_input(placeholder="Ask me anything!")

        if user_query:
            qa_chain = self.setup_qa_chain()
            utils.display_msg(user_query, 'user')

            with st.chat_message("assistant"):
                st_cb = StreamHandler(st.empty())
                response = qa_chain.run(user_query, callbacks=[st_cb])
                st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    obj = CustomDataChatbot()
    obj.main()