LLM Chat App in Python w/ Ollama-py and Streamlit

Description

In this video I walk through the new Ollama Python library, and use it to build a chat app with UI powered by Streamlit. After reviewing some important methods from this library, I touch on Python generators as we construct our chat app, step by step.

Code

# set up python environment
python --version
python -m venv .venv
source .venv/bin/activate
python -m pip install ollama streamlit

# in python shell
# List models
ollama.list()
[model['name'] for model in ollama.list()['models']]

# show info about a model
ollama.show('mistral')
{k: v for k,v in ollama.show('mistral').items() if k != 'license'}

# chat
ollama.chat(
  model='mistral', 
  messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], 
  stream=False
)

import ollama
import streamlit as st

st.title("Ollama Python Chatbot")

# initialize history
if "messages" not in st.session_state:
    st.session_state["messages"] = []

# init models
if "model" not in st.session_state:
    st.session_state["model"] = ""

models = [model["name"] for model in ollama.list()["models"]]
st.session_state["model"] = st.selectbox("Choose your model", models)

def model_res_generator():
    stream = ollama.chat(
        model=st.session_state["model"],
        messages=st.session_state["messages"],
        stream=True,
    )
    for chunk in stream:
        yield chunk["message"]["content"]

# Display chat messages from history on app rerun
for message in st.session_state["messages"]:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("What is up?"):
    # add latest message to history in format {role, content}
    st.session_state["messages"].append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        message = st.write_stream(model_res_generator())
        st.session_state["messages"].append({"role": "assistant", "content": message})