streamlit_app.py

import streamlit as st
import os
import pandas as pd

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
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")

# "./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],
    }
)
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!")

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", 
                                  "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?"}]
    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)

@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
    )

    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):
    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)
            response = generate_llm_response(prompt)
            placeholder = st.empty()
            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)