Tutorial 2: Consumer interfaces with a CKE in a Streamlit chatbot¶
Introduction¶
In this tutorial, you’ll set up a custom retrieval augmented generation (RAG) pipeline to integrate knowledge from a Cortex Knowledge Extension into a chatbot.
This is how it works:
A Streamlit app accepts a prompt from a user.
The prompt is given to the Cortex Search Query API with the configured Cortex Knowledge Extension / Cortex Search Service.
The Streamlit app takes the retrieved documents, puts them into the context window with a custom prompt, and sends it to the Cortex LLM Complete function with a specified LLM.
Note
This tutorial assumes that you have a CKE already available. Go to the Snowflake Marketplace and access one, or use Tutorial 1 to create one.
Step 1. Set up your environment¶
The example below sets up an environment and creates a Streamlit application that you can run in Snowflake to test out a Cortex Knowledge Extension. This assumes the Consumer has access to a Cortex Knowledge Extension that’s been shared by a Provider.
Sign in to Snowsight.
In the left navigation bar, select Projects » Streamlit.
Select + Streamlit App.
The Create Streamlit App window opens.
Enter a name for your app.
In the App location dropdown, select the database and schema for your app.
In the Warehouse dropdown, select the warehouse where you want to run your app and execute queries.
Select Create.
The Streamlit in Snowflake editor opens an example Streamlit app in Viewer mode. Viewer mode allows you to see how the Streamlit application appears to users.
Verify that the correct packages and versions are installed as in the image below.
Step 2: Create a Streamlit app for your CKE chat tester¶
The code below is a simple Streamlit app that allows you to test the CKE. The app uses the Snowflake ML Python package to call the Cortex Knowledge Extension and the Snowflake LLM Complete function. The app allows you to select a Cortex Knowledge Extension, enter a question, and receive a response from the LLM. The app also provides options for debugging and using chat history.
In the navigation menu, select Projects » Streamlit.
Select the Streamlit app you created in the previous step.
In the Streamlit in Snowflake editor, select Edit » Edit code.
The Streamlit in Snowflake editor opens in Edit mode.
In the left navigation bar, select streamlit_app.py to open the code editor.
In the code editor, delete the existing code.
Copy the code below and paste it into the code editor, then select Save » Save and run.
The Streamlit in Snowflake editor runs the app and opens it in Viewer mode.
import streamlit as st
from snowflake.core import Root
from snowflake.cortex import Complete
from snowflake.snowpark.context import get_active_session
MODELS = [
"llama3.1-8b",
"llama3.1-70b",
"llama3.1-405b"
]
def init_messages():
"""Initialize session state messages if not present or if we need to clear."""
if st.session_state.get("clear_conversation") or "messages" not in st.session_state:
st.session_state.messages = []
st.session_state.clear_conversation = False
def init_service_metadata():
"""Load or refresh cortex search services from Snowflake."""
services = session.sql("SHOW CORTEX SEARCH SERVICES IN ACCOUNT;").collect()
service_metadata = []
if services:
for s in services:
svc_name = s["name"]
svc_schema = s["schema_name"]
svc_db = s["database_name"]
svc_search_col = session.sql(
f"DESC CORTEX SEARCH SERVICE {svc_db}.{svc_schema}.{svc_name};"
).collect()[0]["search_column"]
service_metadata.append(
{
"name": svc_name,
"search_column": svc_search_col,
"db": svc_db,
"schema": svc_schema,
}
)
st.session_state.service_metadata = service_metadata
# Initialize selected_cortex_search_service if it doesn't exist
if "selected_cortex_search_service" not in st.session_state and service_metadata:
st.session_state.selected_cortex_search_service = service_metadata[0]["name"]
selected_entry = st.session_state.get("selected_cortex_search_service")
if selected_entry:
# Find matching service metadata
selected_service_metadata = next(
(svc for svc in st.session_state.service_metadata if svc["name"] == selected_entry),
None
)
if selected_service_metadata:
# Store them in session_state
st.session_state.selected_schema = selected_service_metadata["schema"]
st.session_state.selected_db = selected_service_metadata["db"]
elif st.session_state.get("debug", False):
st.write("No matching service found for:", selected_entry)
def init_config_options():
if "service_metadata" not in st.session_state or not st.session_state.service_metadata:
st.sidebar.warning("No Cortex Knowledge Extensions available")
return
st.sidebar.selectbox(
"Select Cortex Knowledge Extension",
[s["name"] for s in st.session_state.service_metadata],
key="selected_cortex_search_service",
)
if st.sidebar.button("Clear conversation"):
st.session_state.clear_conversation = True
# If st.sidebar.toggle isn't available, use st.sidebar.checkbox:
st.sidebar.checkbox("Debug", key="debug", value=False)
st.sidebar.checkbox("Use chat history", key="use_chat_history", value=True)
with st.sidebar.expander("Advanced options"):
st.selectbox("Select model:", MODELS, key="model_name")
st.number_input(
"Select number of context chunks",
value=5,
key="num_retrieved_chunks",
min_value=1,
max_value=10,
)
st.number_input(
"Select number of messages to use in chat history",
value=5,
key="num_chat_messages",
min_value=1,
max_value=10,
)
st.sidebar.expander("Session State").write(st.session_state)
def get_chat_history():
"""Get the last N messages from session state."""
start_index = max(
0, len(st.session_state.messages) - st.session_state.num_chat_messages
)
return st.session_state.messages[start_index : len(st.session_state.messages) - 1]
def complete(model, prompt):
"""Use the chosen Snowflake cortex model to complete a prompt."""
return Complete(model=model, prompt=prompt).replace("$", "\\$")
def make_chat_history_summary(chat_history, question):
"""
Summarize the chat history plus the question using your LLM,
to refine the final search query.
"""
prompt = f"""
[INST]
Based on the chat history below and the question, generate a query that extend the question
with the chat history provided. The query should be in natural language.
Answer with only the query. Do not add any explanation.
<chat_history>
{chat_history}
</chat_history>
<question>
{question}
</question>
[/INST]
"""
summary = complete(st.session_state.model_name, prompt)
if st.session_state.debug:
st.sidebar.text_area("Chat history summary", summary.replace("$", "\\$"), height=150)
return summary
def query_cortex_search_service(query, columns=[], filter={}):
"""
Query the selected cortex search service with the given query and retrieve context documents.
"""
# Safely retrieve from session_state
db = st.session_state.get("selected_db")
schema = st.session_state.get("selected_schema")
if st.session_state.get("debug", False):
st.sidebar.write("Query:", query)
st.sidebar.write("DB:", db)
st.sidebar.write("Schema:", schema)
st.sidebar.write("Service:", st.session_state.selected_cortex_search_service)
cortex_search_service = (
root.databases[db]
.schemas[schema]
.cortex_search_services[st.session_state.selected_cortex_search_service]
)
context_documents = cortex_search_service.search(
query,
columns=columns,
filter=filter,
limit=st.session_state.num_retrieved_chunks
)
results = context_documents.results
if st.session_state.get("debug", False):
st.sidebar.write("Search Results:", results)
service_metadata = st.session_state.service_metadata
search_col = [
s["search_column"] for s in service_metadata
if s["name"] == st.session_state.selected_cortex_search_service
][0].lower()
# Build a context string for the prompt
context_str = ""
context_str_template = (
"Source: {source_url}\n"
"Source ID: {id}\n"
"Excerpt: {chunk}\n\n\n"
)
for i, r in enumerate(results):
context_str += context_str_template.format(
id=i+1,
chunk=r[search_col],
source_url=r["source_url"],
title=r["document_title"],
)
if st.session_state.debug:
st.sidebar.text_area("Context documents", context_str, height=500)
return context_str, results
def create_prompt(user_question):
"""
Combine user question, context from the search service, and chat history
to create a final prompt for the LLM.
"""
if st.session_state.use_chat_history:
chat_history = get_chat_history()
if chat_history != []:
question_summary = make_chat_history_summary(chat_history, user_question)
prompt_context, results = query_cortex_search_service(
question_summary, columns=["chunk", "source_url", "document_title"]
)
else:
prompt_context, results = query_cortex_search_service(
user_question, columns=["chunk", "source_url", "document_title"]
)
else:
prompt_context, results = query_cortex_search_service(
user_question, columns=["chunk", "source_url", "document_title"]
)
chat_history = ""
prompt = f"""
You are a helpful AI assistant with RAG capabilities. When a user asks you a question, you will also be given excerpts from relevant documentation to help answer the question accurately. Please use the context provided and cite your sources using the citation format provided.
Context from documentation:
{prompt_context}
User question:
{user_question}
OUTPUT:
"""
# Add prompt to debug window
if st.session_state.get("debug", False):
st.sidebar.text_area("Complete Prompt", prompt, height=300)
return prompt, results
def post_process_citations(generated_response, results):
"""
Replace {{.StartCitation}}X{{.EndCitation}} with bracketed references to actual product links.
NOTE: If the model references chunks out of range (like 4 if only 2 exist),
consider adding logic to remap or drop invalid references.
"""
used_results = set()
for i, ref in enumerate(results):
old_str = f"{{.StartCitation}}{i+1}{{.EndCitation}}"
replacement = f"[{i+1}]{ref['source_url']})"
new_resp = generated_response.replace(old_str, replacement)
if new_resp != generated_response:
used_results.add(i)
generated_response = new_resp
return generated_response, used_results
# ------------------------------------------------------------------------------
# (2) Main Application (with improved UI)
# ------------------------------------------------------------------------------
def main():
# Optional: wide layout, custom page title
st.set_page_config(
page_title="Cortex Knowledge Extension Chat Tester",
layout="wide",
)
# Optional: a bit of custom CSS for bubble spacing
custom_css = """
<style>
[data-testid="stChatMessage"] {
border-radius: 8px;
margin-bottom: 1rem;
padding: 10px;
}
</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)
# Title or subheader for your app
st.subheader("Cortex Knowledge Extension Chat Tester")
# Initialize metadata and config
init_service_metadata()
init_config_options()
init_messages()
# Icons for user/assistant
icons = {"assistant": "❄️", "user": "👤"}
# Display chat messages from history on app rerun
for message in st.session_state.messages:
with st.chat_message(message["role"], avatar=icons[message["role"]]):
st.markdown(message["content"])
# If there are no services, disable chat
disable_chat = (
"service_metadata" not in st.session_state
or len(st.session_state.service_metadata) == 0
)
# Chat input
if question := st.chat_input("Ask a question...", disabled=disable_chat):
# 1. Store user message
st.session_state.messages.append({"role": "user", "content": question})
# 2. Display user bubble
with st.chat_message("user", avatar=icons["user"]):
st.markdown(question.replace("$", "\\$"))
# 3. Prepare assistant response
with st.chat_message("assistant", avatar=icons["assistant"]):
message_placeholder = st.empty()
# Clean the question
question_safe = question.replace("'", "")
# Build prompt and retrieve docs
prompt, results = create_prompt(question_safe)
with st.spinner("Thinking..."):
generated_response = complete(st.session_state.model_name, prompt)
# Post-process citations
post_processed_response, used_results = post_process_citations(generated_response, results)
# Build references table (only if there are results)
if results:
markdown_table = "\n\n###### References \n\n| Index | Title | Source |\n|------|-------|--------|\n"
for i, ref in enumerate(results):
# Include all references that were found
markdown_table += (
f"| {i+1} | {ref.get('document_title', 'N/A')} | "
f"{ref.get('source_url', 'N/A')} |\n"
)
else:
markdown_table = "\n\n*No references found*"
# Show final assistant message (with references)
message_placeholder.markdown(post_processed_response + markdown_table)
# 4. Append final assistant message to chat history
st.session_state.messages.append(
{"role": "assistant", "content": post_processed_response + markdown_table}
)
# ------------------------------------------------------------------------------
# (3) Entry Point
# ------------------------------------------------------------------------------
if __name__ == "__main__":
session = get_active_session()
root = Root(session)
main()
Step 3: Test the app¶
Click Run to launch the Streamlit application.
Select a CKE from the drop down menu on the left pane under Select Cortex Knowledge Extension.`
Ask a question in the chat text box.
