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.

Code snippet​

For a complete working Reflex example (UI + state management), see the implementation in this repository under reflex/app/pages/genie.py and reflex/app/states/genie_state.py.

app.py
import reflex as rx
from databricks.sdk import WorkspaceClient
import pandas as pd

w = WorkspaceClient()

genie_space_id = "01f0023d28a71e599b5a62f4117916d4"


def display_message(message):
    if "content" in message:
        rx.markdown(message["content"])
    if "data" in message:
        rx.data_table(message["data"])
    if "code" in message:
        rx.accordion.root(
            rx.accordion.item(
                rx.accordion.header(
                    rx.accordion.trigger("Show generated code")
                ),
                rx.accordion.content(
                    rx.code_block(message["code"], language="sql")
                ),
                value="item-1",
            ),
            type="single",
            collapsible=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 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)


class GenieState(rx.State):
    conversation_id: str = ""
    prompt: str = ""

    @rx.event
    async def send_message(self):
        if self.prompt:
            user_prompt = self.prompt
            self.prompt = ""
            # Display user message
            # Then process with assistant
            if self.conversation_id:
                conversation = w.genie.create_message_and_wait(
                    genie_space_id, self.conversation_id, user_prompt
                )
                process_genie_response(conversation)
            else:
                conversation = w.genie.start_conversation_and_wait(genie_space_id, user_prompt)
                self.conversation_id = conversation.conversation_id
                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
reflex
databricks-sdk
pandas