Skip to main content

Chat with a Genie Space

This app uses the AI/BI Genie Conversations API to let users ask questions about your data for instant insights (answers and table-like output). You are also able to collect their feedback on the responses. Visualizations aren't yet supported in the API.

Code snippet​

Refer to the Streamlit Cookbook Genie source code for the full implementation.

app.py
import streamlit as st
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.dashboards import GenieFeedbackRating
import pandas as pd

w = WorkspaceClient()

genie_space_id = "01efe16a65e21836acefb797ae6a8fe4"


def display_message(message):
    if "content" in message:
        st.markdown(message["content"])
    if "data" in message:
        st.dataframe(message["data"])
    if "code" in message:
        with st.expander("Show generated code"):
            st.code(message["code"], language="sql", wrap_lines=True)


def get_query_result(statement_id):
    # For simplicity, let's say data fits in one chunk, query.manifest.total_chunk_count = 1

    result = w.statement_execution.get_statement(statement_id)
    return pd.DataFrame(
        result.result.data_array, columns=[i.name for i in result.manifest.schema.columns]
    )


def collect_feedback(message_id: str):
    rating = st.feedback("thumbs", key=f"feedback_{message_id}")
    mapping = {1: GenieFeedbackRating.POSITIVE, 0: GenieFeedbackRating.NEGATIVE}
    if rating and message["message_id"]:
        w.genie.send_message_feedback(
            genie_space_id, st.session_state.conversation_id, message["message_id"], mapping[rating]
        )


def process_genie_response(response):
    for i in response.attachments:
        if i.text:
            message = {"role": "assistant", "content": i.text.content}
            display_message(message)
        elif i.query:
            data = get_query_result(response.query_result.statement_id)
            message = {
                "role": "assistant", "content": i.query.description, "data": data, "code": i.query.query
            }
            display_message(message)
    collect_feedback(response.message_id)


if prompt := st.chat_input("Ask your question..."):
    # Refer to actual app code for chat history persistence on rerun

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

    with st.chat_message("assistant"):
        if st.session_state.get("conversation_id"):
            conversation = w.genie.create_message_and_wait(
                genie_space_id, st.session_state.conversation_id, prompt
            )
            process_genie_response(conversation)
        else:
            conversation = w.genie.start_conversation_and_wait(genie_space_id, prompt)
            process_genie_response(conversation)
info

Copy and paste the Genie space ID from the Genie UI URL as rooms/SPACE-ID?o=.

Resources​

Permissions​

Your app service principal needs the following permissions:

  • SELECT on the Unity Catalog table
  • CAN USE the SQL warehouse
  • CAN VIEW the Genie Space

Dependencies​

requirements.txt
streamlit
databricks-sdk
pandas