Chat Application using Streamlit and Text Bison

I am a Java backend engineer and dealing with the different front-end frameworks to create a passable user interface is often a challenging task for me. On and off I have wrangled with different javascript frameworks, CSS frameworks, but I have never felt confident enough to call myself a full-stack developer.

Streamlit is great for developers like me who would love to create a UI but lack the front-end chops.

So my objective in this post is to go over how I created a small chatbot application using Streamlit and Google’s text-bison as the base Generative AI model.

Basic Structure

To start with Streamlit, I first wanted to have a basic structure that has all the elements of a chat application but with a canned response. This is very simple using streamlit.

Assuming that streamlit is installed using:

pip install streamlit

a basic python code that prompts a user and responds with a canned response looks like this:

import streamlit as st

prompt: str = st.chat_input("Enter a prompt here")

USER = "user"
ASSISTANT = "assistant"

if prompt:
    st.chat_message(USER).write(prompt)
    st.chat_message(ASSISTANT).write(f"You wrote {prompt}")

Running this python code using:

streamlit run app.py

a good UI that looks like this shows up:

It already looks great and responds to a user prompt by returning canned content as a response.

Adding in a second prompt, however, overwrites the first set of results:

The reason this happens is because Streamlit re-runs the entire script when a user interacts with a widget, the prompt input in this instance.

So the question is how can the existing interaction be retained. This is where session data comes in, the content of the chat session can be maintained using a session state and the updated code looks like this:

import streamlit as st
from dataclasses import dataclass


@dataclass
class Message:
    actor: str
    payload: str


USER = "user"
ASSISTANT = "ai"
MESSAGES = "messages"
if MESSAGES not in st.session_state:
    st.session_state[MESSAGES] = [Message(actor=ASSISTANT, payload="Hi!How can I help you?")]

msg: Message
for msg in st.session_state[MESSAGES]:
    st.chat_message(msg.actor).write(msg.payload)

prompt: str = st.chat_input("Enter a prompt here")

if prompt:
    st.session_state[MESSAGES].append(Message(actor=USER, payload=prompt))
    st.chat_message(USER).write(prompt)
    response: str = f"You wrote {prompt}"
    st.session_state[MESSAGES].append(Message(actor=ASSISTANT, payload=response))
    st.chat_message(ASSISTANT).write(response) 

As the chat progresses, the content is maintained in a session state variable called MESSAGES, on new prompts, existing content from the messages is printed to the screen and then the prompt is processed, this way the existing messages are not lost.

Wiring in Text Bison

Text Bison is Google’s foundational LLM model for natural language text generation. Wiring in text bison is simple, first pull in the google-cloud-aiplatform as a dependency:

pip install google-cloud-aiplatform

and the following code wires in text-bison and uses it now to answer the prompts!

from dataclasses import dataclass

import streamlit as st
from vertexai.preview.language_models import TextGenerationModel


@dataclass
class Message:
    actor: str
    payload: str


def get_llm() -> TextGenerationModel:
    return TextGenerationModel.from_pretrained("text-bison@001")


USER = "user"
ASSISTANT = "ai"
MESSAGES = "messages"
if MESSAGES not in st.session_state:
    st.session_state[MESSAGES] = [Message(actor=ASSISTANT, payload="Hi!How can I help you?")]

msg: Message
for msg in st.session_state[MESSAGES]:
    st.chat_message(msg.actor).write(msg.payload)

prompt: str = st.chat_input("Enter a prompt here")

if prompt:
    st.session_state[MESSAGES].append(Message(actor=USER, payload=prompt))
    st.chat_message(USER).write(prompt)
    response: str = get_llm().predict(prompt=prompt).text
    st.session_state[MESSAGES].append(Message(actor=ASSISTANT, payload=response))
    st.chat_message(ASSISTANT).write(response)

This works great! I can ask some basic questions and get answers already:

However, one problem with this is that the context of the conversation does not carry over. Each question is independent of the previous one. See for eg:

So, now how can the context of a conversation be maintained?

This is where something like langchain should be introduced. Langchain provides ways to hold and pass on conversation histories as documented here, using this the code looks like this:

First to install it:

pip install langchain

and using it in the code:

from dataclasses import dataclass

import streamlit as st
from langchain.chains import LLMChain
from langchain.llms import VertexAI
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate


@dataclass
class Message:
    actor: str
    payload: str


@st.cache_resource
def get_llm() -> VertexAI:
    return VertexAI(model_name="text-bison@001")


def get_llm_chain():
    template = """You are a nice chatbot having a conversation with a human.

    Previous conversation:
    {chat_history}

    New human question: {question}
    Response:"""
    prompt_template = PromptTemplate.from_template(template)
    # Notice that we need to align the `memory_key`
    memory = ConversationBufferMemory(memory_key="chat_history")
    conversation = LLMChain(
        llm=get_llm(),
        prompt=prompt_template,
        verbose=True,
        memory=memory
    )
    return conversation


USER = "user"
ASSISTANT = "ai"
MESSAGES = "messages"


def initialize_session_state():
    if MESSAGES not in st.session_state:
        st.session_state[MESSAGES] = [Message(actor=ASSISTANT, payload="Hi!How can I help you?")]
    if "llm_chain" not in st.session_state:
        st.session_state["llm_chain"] = get_llm_chain()


def get_llm_chain_from_session() -> LLMChain:
    return st.session_state["llm_chain"]


initialize_session_state()

msg: Message
for msg in st.session_state[MESSAGES]:
    st.chat_message(msg.actor).write(msg.payload)

prompt: str = st.chat_input("Enter a prompt here")

if prompt:
    st.session_state[MESSAGES].append(Message(actor=USER, payload=prompt))
    st.chat_message(USER).write(prompt)

    with st.spinner("Please wait.."):
        llm_chain = get_llm_chain_from_session()
        response: str = llm_chain({"question": prompt})["text"]
        st.session_state[MESSAGES].append(Message(actor=ASSISTANT, payload=response))
        st.chat_message(ASSISTANT).write(response)

This likely is a little more complicated than it needs to be, chalk it down to me being a beginner Python programmer, the gist of it is that an LLM chain has been created which wires in the VertexAI text-bison model, a custom prompt template that considers the previous conversation history and something to hold the history.

A conversation now carries the context from one prompt to the next:

Conclusion

This concludes a quick tour of Streamlit and using it to create a passable chat interface. I am amazed by how simple it has been to create this application without needing to muddle through CSS and javascript frameworks.

The entire sample is available in my GitHub repository here — https://github.com/bijukunjummen/genai-chat