to implement rag in the streamlit chatbot app

archive/test.py

@@ -0,0 +1,62 @@
+import utils
+import os
+
+import openai
+from llama_index import SimpleDirectoryReader
+from llama_index import Document
+from llama_index import VectorStoreIndex
+from llama_index import ServiceContext
+from llama_index.llms import OpenAI
+
+from llama_index.embeddings import HuggingFaceEmbedding
+from trulens_eval import Tru
+from llama_index.memory import ChatMemoryBuffer
+
+from utils import get_prebuilt_trulens_recorder
+import time
+
+openai.api_key = utils.get_openai_api_key()
+
+documents = SimpleDirectoryReader(
+    input_files=["./raw_documents/HI_Knowledge_Base.pdf"]
+).load_data()
+
+document = Document(text="\n\n".join([doc.text for doc in documents]))
+
+### gpt-4-1106-preview
+### gpt-3.5-turbo-1106 / gpt-3.5-turbo
+print("Initializing GPT 3.5 ..")
+llm = OpenAI(model="gpt-3.5-turbo-1106", temperature=0.1)
+
+print("Initializing bge-small-en-v1.5 embedding model ..")
+embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
+
+print("Creating vector store ..")
+service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)
+index = VectorStoreIndex.from_documents([document], service_context=service_context)
+
+if False:
+    query_engine = index.as_query_engine(streaming=True)
+else:
+    memory = ChatMemoryBuffer.from_defaults(token_limit=15000)
+
+    # chat_engine = index.as_query_engine(streaming=True)
+    chat_engine = index.as_chat_engine(
+        chat_mode="context",
+        memory=memory
+    )
+
+while True:
+    input_str = input("[User]: ")
+    if input_str == "END":
+        break
+    
+    # res = chat_engine.query(input_str)
+    res = chat_engine.stream_chat(input_str)
+
+    bot_response = ""
+    print("[Bot]: ", end="")
+    for s in res.response_gen:
+        bot_response += s
+        print(s, end="")
+    print("")

streamlit_app.py

@@ -3,7 +3,16 @@ import os
 import pandas as pd
 
 import openai
-from openai import OpenAI
+
+# from openai import OpenAI
+from llama_index.llms import OpenAI
+
+from llama_index import SimpleDirectoryReader
+from llama_index import Document
+from llama_index import VectorStoreIndex
+from llama_index import ServiceContext
+from llama_index.embeddings import HuggingFaceEmbedding
+from llama_index.memory import ChatMemoryBuffer
 
 import pkg_resources
 import shutil
@@ -35,6 +44,14 @@ shutil.copytree(ux_path, "./ux")
 st.set_page_config(page_title="💬 Open AI Chatbot")
 openai_api = os.getenv("OPENAI_API_KEY")
 
+# "./raw_documents/HI_Knowledge_Base.pdf"
+input_files = ["./raw_documents/HI Chapter Summary Version 1.3.pdf"]
+embedding_model = "BAAI/bge-small-en-v1.5"
+system_content = ("You are a helpful study assistant. "
+                  "You do not respond as 'User' or pretend to be 'User'. "
+                  "You only respond once as 'Assistant'."
+)
+
 data_df = pd.DataFrame(
     {
         "Completion": [30, 40, 100, 10],
@@ -83,6 +100,9 @@ with st.sidebar:
 
     st.markdown("📖 Reach out to SakiMilo to learn how to create this app!")
 
+if "init" not in st.session_state.keys():
+    st.session_state.init = {"warm_start": "No"}
+
 # Store LLM generated responses
 if "messages" not in st.session_state.keys():
     st.session_state.messages = [{"role": "assistant", 
@@ -96,23 +116,66 @@ for message in st.session_state.messages:
 def clear_chat_history():
     st.session_state.messages = [{"role": "assistant", 
                                   "content": "How may I assist you today?"}]
+    chat_engine = get_query_engine(input_files=input_files, 
+                                   llm_model=selected_model, 
+                                   temperature=temperature,
+                                   embedding_model=embedding_model,
+                                   system_content=system_content)
+    chat_engine.reset()
 st.sidebar.button("Clear Chat History", on_click=clear_chat_history)
 
-def generate_llm_response(client, prompt_input):
-    system_content = ("You are a helpful assistant. "
-                      "You do not respond as 'User' or pretend to be 'User'. "
-                      "You only respond once as 'Assistant'."
+@st.cache_resource
+def get_document_object(input_files):
+    documents = SimpleDirectoryReader(input_files=input_files).load_data()
+    document = Document(text="\n\n".join([doc.text for doc in documents]))
+    return document
+
+@st.cache_resource
+def get_llm_object(selected_model, temperature):
+    llm = OpenAI(model=selected_model, temperature=temperature)
+    return llm
+
+@st.cache_resource
+def get_embedding_model(model_name):
+    embed_model = HuggingFaceEmbedding(model_name=model_name)
+    return embed_model
+
+@st.cache_resource
+def get_query_engine(input_files, llm_model, temperature, 
+                     embedding_model, system_content):
+
+    document = get_document_object(input_files)
+    llm = get_llm_object(llm_model, temperature)
+    embedded_model = get_embedding_model(embedding_model)
+
+    service_context = ServiceContext.from_defaults(llm=llm, embed_model=embedded_model)
+    index = VectorStoreIndex.from_documents([document], service_context=service_context)
+    memory = ChatMemoryBuffer.from_defaults(token_limit=15000)
+
+    # chat_engine = index.as_query_engine(streaming=True)
+    chat_engine = index.as_chat_engine(
+        chat_mode="context",
+        memory=memory,
+        system_prompt=system_content
     )
 
-    completion = client.chat.completions.create(
-                    model=selected_model,
-                    messages=[
-                            {"role": "system", "content": system_content},
-                    ] + st.session_state.messages,
-                    temperature=temperature,
-                    stream=True
-    )
-    return completion
+    return chat_engine
+
+def generate_llm_response(prompt_input):
+    chat_engine = get_query_engine(input_files=input_files, 
+                                   llm_model=selected_model, 
+                                   temperature=temperature,
+                                   embedding_model=embedding_model,
+                                   system_content=system_content)
+    
+    # st.session_state.messages
+    response = chat_engine.stream_chat(prompt_input)
+    return response
+
+# Warm start
+if st.session_state.init["warm_start"] == "No":
+    clear_chat_history()
+    st.session_state.init["warm_start"] = "Yes"
 
 # User-provided prompt
 if prompt := st.chat_input(disabled=not openai_api):
@@ -125,13 +188,13 @@ if prompt := st.chat_input(disabled=not openai_api):
 if st.session_state.messages[-1]["role"] != "assistant":
     with st.chat_message("assistant"):
         with st.spinner("Thinking..."):
-            response = generate_llm_response(client, prompt)
+            # response = generate_llm_response(client, prompt)
+            response = generate_llm_response(prompt)
             placeholder = st.empty()
             full_response = ""
-            for chunk in response:
-                if chunk.choices[0].delta.content is not None:
-                    full_response += chunk.choices[0].delta.content
-                    placeholder.markdown(full_response)
+            for token in response.response_gen:
+                full_response += token
+                placeholder.markdown(full_response)
             placeholder.markdown(full_response)
     message = {"role": "assistant", "content": full_response}
     st.session_state.messages.append(message)

streamlit_app_archive.py

@@ -0,0 +1,137 @@
+import streamlit as st
+import os
+import pandas as pd
+
+import openai
+from openai import OpenAI
+
+import pkg_resources
+import shutil
+import main
+
+### To trigger trulens evaluation
+main.main()
+
+### Finally, start streamlit app
+leaderboard_path = pkg_resources.resource_filename(
+                    "trulens_eval", "Leaderboard.py"
+                )
+evaluation_path = pkg_resources.resource_filename(
+                    "trulens_eval", "pages/Evaluations.py"
+                )
+ux_path = pkg_resources.resource_filename(
+                    "trulens_eval", "ux"
+                )
+
+os.makedirs("./pages", exist_ok=True)
+shutil.copyfile(leaderboard_path, os.path.join("./pages", "1_Leaderboard.py"))
+shutil.copyfile(evaluation_path, os.path.join("./pages", "2_Evaluations.py"))
+
+if os.path.exists("./ux"):
+    shutil.rmtree("./ux")
+shutil.copytree(ux_path, "./ux")
+
+# App title
+st.set_page_config(page_title="💬 Open AI Chatbot")
+openai_api = os.getenv("OPENAI_API_KEY")
+
+data_df = pd.DataFrame(
+    {
+        "Completion": [30, 40, 100, 10],
+    }
+)
+data_df.index = ["Chapter 1", "Chapter 2", "Chapter 3", "Chapter 4"]
+
+# Replicate Credentials
+with st.sidebar:
+    st.title("💬 Open AI Chatbot")
+    st.write("This chatbot is created using the GPT model from Open AI.")
+    if openai_api:
+        pass
+    elif "OPENAI_API_KEY" in st.secrets:
+        st.success("API key already provided!", icon="✅")
+        openai_api = st.secrets["OPENAI_API_KEY"]
+    else:
+        openai_api = st.text_input("Enter OpenAI API token:", type="password")
+        if not (openai_api.startswith("sk-") and len(openai_api)==51):
+            st.warning("Please enter your credentials!", icon="⚠️")
+        else:
+            st.success("Proceed to entering your prompt message!", icon="👉")
+
+    ### for streamlit purpose
+    os.environ["OPENAI_API_KEY"] = openai_api
+
+    st.subheader("Models and parameters")
+    selected_model = st.sidebar.selectbox("Choose an OpenAI model", 
+                                          ["gpt-3.5-turbo-1106", "gpt-4-1106-preview"], 
+                                           key="selected_model")
+    temperature = st.sidebar.slider("temperature", min_value=0.01, max_value=2.0, 
+                                    value=0.1, step=0.01)
+    st.data_editor(
+        data_df,
+        column_config={
+            "Completion": st.column_config.ProgressColumn(
+                            "Completion %",
+                            help="Percentage of content covered",
+                            format="%.1f%%",
+                            min_value=0,
+                            max_value=100,
+            ),
+        },
+        hide_index=False,
+    )
+
+    st.markdown("📖 Reach out to SakiMilo to learn how to create this app!")
+
+# Store LLM generated responses
+if "messages" not in st.session_state.keys():
+    st.session_state.messages = [{"role": "assistant", 
+                                  "content": "How may I assist you today?"}]
+
+# Display or clear chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
+
+def clear_chat_history():
+    st.session_state.messages = [{"role": "assistant", 
+                                  "content": "How may I assist you today?"}]
+st.sidebar.button("Clear Chat History", on_click=clear_chat_history)
+
+def generate_llm_response(client, prompt_input):
+    system_content = ("You are a helpful assistant. "
+                      "You do not respond as 'User' or pretend to be 'User'. "
+                      "You only respond once as 'Assistant'."
+    )
+
+    completion = client.chat.completions.create(
+                    model=selected_model,
+                    messages=[
+                            {"role": "system", "content": system_content},
+                    ] + st.session_state.messages,
+                    temperature=temperature,
+                    stream=True
+    )
+    return completion
+
+# User-provided prompt
+if prompt := st.chat_input(disabled=not openai_api):
+    client = OpenAI()
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.write(prompt)
+
+# Generate a new response if last message is not from assistant
+if st.session_state.messages[-1]["role"] != "assistant":
+    with st.chat_message("assistant"):
+        with st.spinner("Thinking..."):
+            response = generate_llm_response(client, prompt)
+            placeholder = st.empty()
+            full_response = ""
+            for chunk in response:
+                if chunk.choices[0].delta.content is not None:
+                    full_response += chunk.choices[0].delta.content
+                    placeholder.markdown(full_response)
+            placeholder.markdown(full_response)
+    message = {"role": "assistant", "content": full_response}
+    st.session_state.messages.append(message)