Updated source code with genie room as added resource (#574)

[updated source code with genie room as added
resource](912fd4f)

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: vivian-xie-db

conversational-agent-app/README.md

@@ -15,10 +15,7 @@ date: 2025-06-09
 This repository demonstrates how to integrate Databricks' AI/BI Genie Conversation APIs into custom Databricks Apps applications, allowing users to interact with their structured data using natural language.
 
 You can also click the Generate insights button and generate deep analysis and trends of your data.
-![](./assets/insights1.png)
-![](./assets/insights2.png)
-
-
+![](./assets/genie2.png)
 
 ## Overview
 
@@ -29,6 +26,7 @@ The Databricks Genie Conversation APIs (in Public Preview) enable you to embed A
 - Get SQL-powered insights without writing code
 - Follow up with contextual questions in a conversation thread
 
+
 ## Key Features
 
 - **Powered by Databricks Apps**: Deploy and run directly from your Databricks workspace with built-in security and scaling
@@ -58,6 +56,9 @@ The app can be installed through Databricks Marketplace. If you prefer to clone
    - Go to the Apps tab and click the **Create app** button. Fill in the necessary fields and click **Next: Configuration**
       - To reuse an existing app, click the link to your app in the **Name** column to go to the detail page of the app, then click **Edit**
    - In the App resources section, click **+ Add resource** and select **Serving endpoint**. Choose a chat endpoint, grant **CAN_QUERY** permission and name it 'serving_endpoint'.
+   select **Genie Space**. Choose a genie space, grant **CAN_RUN** permission and name it 'genie_space'.
+   ![](./assets/genie1.png)
+
 
 For detailed instructions on configuring resources, refer to the [official Databricks documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/resources#configure-resources-for-your-app).
 
@@ -70,7 +71,7 @@ For detailed instructions on configuring resources, refer to the [official Datab
 For more details, refer to the [official Databricks documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/deploy).
 
 ## Troubleshooting
-1. After you install the app from Markeplace, check the Authorization page for API scope details. 
+1. After you install the app from Marketplace, check the Authorization page for API scope details. 
 ![](./assets/genie-space12.png)
 
 Then you open the url link the first time, ensure that you see this OBO scope authorization page, which has all four scopes:

conversational-agent-app/app.py

@@ -22,6 +22,9 @@
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Get the Genie space ID from environment variable
+GENIE_SPACE_ID = os.environ.get("GENIE_SPACE")
+
 # Create Dash app
 app = dash.Dash(
     __name__,
@@ -58,21 +61,7 @@
 # Define the layout
 app.layout = html.Div([
     html.Div([
-        dcc.Store(id="selected-space-id", data=None, storage_type="local"),
-        dcc.Store(id="spaces-list", data=[]),
         dcc.Store(id="user-info", data={"initial": "Y", "username": "You"}),
-        # Space selection overlay
-        html.Div([
-            html.Div([
-                html.Div([
-                    html.Span(className="space-select-spinner"),
-                    "Loading Genie Spaces..."
-                ], id="space-select-title", className="space-select-title"),
-                dcc.Dropdown(id="space-dropdown", options=[], placeholder="Choose a Genie Space", className="space-select-dropdown", optionHeight=60, searchable=True),
-                html.Button("Select", id="select-space-button", className="space-select-button"),
-                html.Div(id="space-select-error", className="space-select-error")
-            ], className="space-select-card")
-        ], id="space-select-container", className="space-select-container"),
         # Top navigation bar
         html.Div([
             # Left component containing both nav-left and sidebar
@@ -191,16 +180,16 @@
                     html.Div("Always review the accuracy of responses.", className="disclaimer-fixed")
                 ], id="fixed-input-wrapper", className="fixed-input-wrapper"),
             ], id="chat-container", className="chat-container"),
-        ], id="main-content", className="main-content", style={"display": "none"}),
+        ], id="main-content", className="main-content"),
 
         html.Div(id='dummy-output'),
+        html.Div(id='dummy-insight-scroll'),
         dcc.Store(id="chat-trigger", data={"trigger": False, "message": ""}),
         dcc.Store(id="chat-history-store", data=[]),
         dcc.Store(id="query-running-store", data=False),
         dcc.Store(id="session-store", data={"current_session": None}),
-        html.Div(id='dummy-insight-scroll')
-    ], id="app-inner-layout"),
-], id="root-container")
+    ], id="root-container", className="root-container")
+], id="app-container", className="app-container")
 
 # Store chat history
 chat_history = []
@@ -411,11 +400,10 @@ def handle_all_inputs(s1_clicks, s2_clicks, s3_clicks, s4_clicks, send_clicks, s
      Output("query-running-store", "data", allow_duplicate=True)],
     [Input("chat-trigger", "data")],
     [State("chat-messages", "children"),
-     State("chat-history-store", "data"),
-     State("selected-space-id", "data")],
+     State("chat-history-store", "data")],
     prevent_initial_call=True
 )
-def get_model_response(trigger_data, current_messages, chat_history, selected_space_id):
+def get_model_response(trigger_data, current_messages, chat_history):
     if not trigger_data or not trigger_data.get("trigger"):
         return dash.no_update, dash.no_update, dash.no_update, dash.no_update
 
@@ -425,9 +413,8 @@ def get_model_response(trigger_data, current_messages, chat_history, selected_sp
 
     try:
         headers = request.headers
-        # user_token = os.environ.get("DATABRICKS_TOKEN")
         user_token = headers.get('X-Forwarded-Access-Token')
-        response, query_text = genie_query(user_input, user_token, selected_space_id)
+        response, query_text = genie_query(user_input, user_token, GENIE_SPACE_ID)
 
         if isinstance(response, str):
             # Escape square brackets to prevent markdown auto-linking
@@ -747,84 +734,42 @@ def generate_insights(n_clicks, btn_id, chat_history):
     if df is None:
         return html.Div("No data available for insights.", style={"color": "red"})
     insights = call_llm_for_insights(df)
-    return html.Div(
-        dcc.Markdown(insights),
-        style={"marginTop": "32px", "background": "#f4f4f4", "padding": "16px", "borderRadius": "4px"},
-        className="insight-output"
-    )
+    return html.Div([
+        html.Div([
+            dcc.Markdown(insights, className="insight-content")
+        ], className="insight-body")
+    ], className="insight-wrapper")
+
 
 # Callback to fetch spaces on load
+# Initialize welcome title and description from space info
 @app.callback(
-    Output("spaces-list", "data"),
-    Input("space-select-container", "id"),
+    [Output("welcome-title", "children"),
+     Output("welcome-description", "children")],
+    Input("user-info", "id"),
     prevent_initial_call=False
 )
-def fetch_spaces(_):
+def initialize_welcome_info(_):
+    if not GENIE_SPACE_ID:
+        return DEFAULT_WELCOME_TITLE, DEFAULT_WELCOME_DESCRIPTION
+    
     try:
+        # Try to fetch space information to get title and description
         headers = request.headers
-        # token = os.environ.get("DATABRICKS_TOKEN")
         token = headers.get('X-Forwarded-Access-Token')
         host = os.environ.get("DATABRICKS_HOST")
         client = GenieClient(host=host, space_id="", token=token)
-        spaces = client.list_spaces()
-        return spaces
+        
+        # Get the specific space details
+        space_details = client.get_space(GENIE_SPACE_ID)
+        
+        title = space_details.get("title", DEFAULT_WELCOME_TITLE)
+        description = space_details.get("description", DEFAULT_WELCOME_DESCRIPTION)
+        return title, description
+            
     except Exception as e:
-        return []
-
-# Populate dropdown options
-@app.callback(
-    Output("space-dropdown", "options"),
-    Input("spaces-list", "data"),
-    prevent_initial_call=False
-)
-def update_space_dropdown(spaces):
-    if not spaces:
-        return []
-    options = []
-    for s in spaces:
-        title = s.get('title', '')
-        space_id = s.get('space_id', '')
-        label_lines = [title]
-        label_lines.append(space_id)
-        label = " | ".join(label_lines)  # or use '\\n'.join(label_lines) for multi-line (but most browsers will show as a single line)
-        options.append({"label": label, "value": space_id})
-    return options
-
-# Handle space selection
-@app.callback(
-    [Output("selected-space-id", "data", allow_duplicate=True),
-     Output("space-select-container", "style"),
-     Output("main-content", "style"),
-     Output("space-select-error", "children"),
-     Output("welcome-title", "children"),
-     Output("welcome-description", "children")],
-    Input("select-space-button", "n_clicks"),
-    State("space-dropdown", "value"),
-    State("spaces-list", "data"),
-    prevent_initial_call=True
-)
-def select_space(n_clicks, space_id, spaces):
-    if not n_clicks:
-        return dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update
-    if not space_id:
-        return dash.no_update, {"display": "flex", "flexDirection": "column", "alignItems": "center", "justifyContent": "center", "height": "100vh"}, {"display": "none"}, "Please select a Genie space.", dash.no_update, dash.no_update
-    # Find the selected space's title and description
-    selected = next((s for s in spaces if s["space_id"] == space_id), None)
-    title = selected["title"] if selected and selected.get("title") else DEFAULT_WELCOME_TITLE
-    description = selected["description"] if selected and selected.get("description") else DEFAULT_WELCOME_DESCRIPTION
-    return space_id, {"display": "none"}, {"display": "block"}, "", title, description
-
-# Add a callback to control visibility of main-content and space-select-container
-@app.callback(
-    [Output("main-content", "style", allow_duplicate=True), Output("space-select-container", "style", allow_duplicate=True)],
-    Input("selected-space-id", "data"),
-    prevent_initial_call=True
-)
-def toggle_main_ui(selected_space_id):
-    if selected_space_id:
-        return {"display": "block"}, {"display": "none"}
-    else:
-        return {"display": "none"}, {"display": "flex", "flexDirection": "column", "alignItems": "center", "justifyContent": "center", "height": "100vh"}
+        logger.warning(f"Could not fetch space information: {e}")
+        return DEFAULT_WELCOME_TITLE, DEFAULT_WELCOME_DESCRIPTION
 
 # Add clientside callback for scrolling to bottom of chat when insight is generated
 app.clientside_callback(
@@ -847,39 +792,20 @@ def toggle_main_ui(selected_space_id):
     prevent_initial_call=True
 )
 
-@app.callback(
-    Output("selected-space-id", "data", allow_duplicate=True),
-    Input("logout-button", "n_clicks"),
-    prevent_initial_call=True
-)
-def logout_and_clear_space(n_clicks):
-    if n_clicks:
-        return None
-    return dash.no_update
 
-# Add a callback to control the root-container style to prevent scrolling when overlay is visible
+
+# Add a callback to control the root-container style
 @app.callback(
     Output("root-container", "style"),
-    Input("selected-space-id", "data"),
+    Input("user-info", "id"),
     prevent_initial_call=False
 )
-def set_root_style(selected_space_id):
-    if selected_space_id:
+def set_root_style(_):
+    if GENIE_SPACE_ID:
         return {"height": "auto", "overflow": "auto"}
     else:
         return {"height": "100vh", "overflow": "hidden"}
 
-# Add a callback to update the title based on spaces-list
-@app.callback(
-    Output("space-select-title", "children"),
-    Input("spaces-list", "data"),
-    prevent_initial_call=False
-)
-def update_space_select_title(spaces):
-    if not spaces:
-        return [html.Span(className="space-select-spinner"), "Loading Genie Spaces..."]
-    return "Select a Genie Space"
-
 @app.callback(
     Output("query-tooltip", "className"),
     Input("query-running-store", "data"),

conversational-agent-app/app.yaml

@@ -4,4 +4,6 @@ command:
 
 env:
 - name: "SERVING_ENDPOINT_NAME"
-  valueFrom: "serving_endpoint"
+  valueFrom: "serving-endpoint"
+- name: "GENIE_SPACE"
+  valueFrom: "genie-space"