Simple Functional Chatbot with Streamlit and Cortex API in Snowflake.

chatbot.py

import streamlit as st
import json
from snowflake.snowpark import Session
session = Session.builder.getOrCreate()

st.title("☃️ Frosty")

# Initialize the chat messages history
if "messages" not in st.session_state.keys():
    st.session_state.messages = [{"role": "assistant", "content": "How can I help?"}]

# Prompt for user input and save
if prompt := st.chat_input():
    st.session_state.messages.append({"role": "user", "content": prompt})

# display the existing chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# If last message is not from assistant, we need to generate a new response
if st.session_state.messages[-1]["role"] != "assistant":
    # Call LLM
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            query = json.dumps([{"role": m["role"], "content": m["content"]} for m in st.session_state.messages]) 
            results = session.sql(f"select snowflake.cortex.complete('mistral-large',parse_json($${query}$$))").first()
            response = results[0]
            st.write(response)

    message = {"role": "assistant", "content": response}
    st.session_state.messages.append(message)