rag-chatbot

1
# Import necessary libraries
2
import databutton as db
3
import streamlit as st
4
import openai
5
from brain import get_index_for_pdf
6
from langchain.chains import RetrievalQA
7
from langchain.chat_models import ChatOpenAI
8
import os
9

10
# Set the title for the Streamlit app
11
st.title("RAG enhanced Chatbot")
12

13
# Set up the OpenAI API key from databutton secrets
14
os.environ["OPENAI_API_KEY"] = db.secrets.get("OPENAI_API_KEY")
15
openai.api_key = db.secrets.get("OPENAI_API_KEY")
16

17

18
# Cached function to create a vectordb for the provided PDF files
19
@st.cache_data
20
def create_vectordb(files, filenames):
21
    # Show a spinner while creating the vectordb
22
    with st.spinner("Vector database"):
23
        vectordb = get_index_for_pdf(
24
            [file.getvalue() for file in files], filenames, openai.api_key
25
        )
26
    return vectordb
27

28

29
# Upload PDF files using Streamlit's file uploader
30
pdf_files = st.file_uploader("", type="pdf", accept_multiple_files=True)
31

32
# If PDF files are uploaded, create the vectordb and store it in the session state
33
if pdf_files:
34
    pdf_file_names = [file.name for file in pdf_files]
35
    st.session_state["vectordb"] = create_vectordb(pdf_files, pdf_file_names)
36

37
# Define the template for the chatbot prompt
38
prompt_template = """
39
    You are a helpful Assistant who answers to users questions based on multiple contexts given to you.
40

41
    Keep your answer short and to the point.
42
    
43
    The evidence are the context of the pdf extract with metadata. 
44
    
45
    Carefully focus on the metadata specially 'filename' and 'page' whenever answering.
46
    
47
    Make sure to add filename and page number at the end of sentence you are citing to.
48
        
49
    Reply "Not applicable" if text is irrelevant.
50
     
51
    The PDF content is:
52
    {pdf_extract}
53
"""
54

55
# Get the current prompt from the session state or set a default value
56
prompt = st.session_state.get("prompt", [{"role": "system", "content": "none"}])
57

58
# Display previous chat messages
59
for message in prompt:
60
    if message["role"] != "system":
61
        with st.chat_message(message["role"]):
62
            st.write(message["content"])
63

64
# Get the user's question using Streamlit's chat input
65
question = st.chat_input("Ask anything")
66

67
# Handle the user's question
68
if question:
69
    vectordb = st.session_state.get("vectordb", None)
70
    if not vectordb:
71
        with st.message("assistant"):
72
            st.write("You need to provide a PDF")
73
            st.stop()
74

75
    # Search the vectordb for similar content to the user's question
76
    search_results = vectordb.similarity_search(question, k=3)
77
    # search_results
78
    pdf_extract = "/n ".join([result.page_content for result in search_results])
79

80
    # Update the prompt with the pdf extract
81
    prompt[0] = {
82
        "role": "system",
83
        "content": prompt_template.format(pdf_extract=pdf_extract),
84
    }
85

86
    # Add the user's question to the prompt and display it
87
    prompt.append({"role": "user", "content": question})
88
    with st.chat_message("user"):
89
        st.write(question)
90

91
    # Display an empty assistant message while waiting for the response
92
    with st.chat_message("assistant"):
93
        botmsg = st.empty()
94

95
    # Call ChatGPT with streaming and display the response as it comes
96
    response = []
97
    result = ""
98
    for chunk in openai.ChatCompletion.create(
99
        model="gpt-3.5-turbo", messages=prompt, stream=True
100
    ):
101
        text = chunk.choices[0].get("delta", {}).get("content")
102
        if text is not None:
103
            response.append(text)
104
            result = "".join(response).strip()
105
            botmsg.write(result)
106

107
    # Add the assistant's response to the prompt
108
    prompt.append({"role": "assistant", "content": result})
109

110
    # Store the updated prompt in the session state
111
    st.session_state["prompt"] = prompt
112
    prompt.append({"role": "assistant", "content": result})
113

114
    # Store the updated prompt in the session state
115
    st.session_state["prompt"] = prompt
116