Fix Genie API code snippet

Author: pbv0

streamlit/views/genie_api.py

@@ -1,10 +1,10 @@
-import streamlit as st
+from typing import Dict
+
+import pandas as pd
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors.sdk import OperationFailed
 from databricks.sdk.service.dashboards import GenieMessage
-import pandas as pd
-from typing import Dict
 
+import streamlit as st
 
 w = WorkspaceClient()
 
@@ -16,23 +16,24 @@
     to let users ask questions about your data for instant insights.
     """
 )
-st.warning("Public Preview")
 
 tab_a, tab_b, tab_c = st.tabs(["**Try it**", "**Code snippet**", "**Requirements**"])
 
 with tab_a:
+
     def reset_conversation():
         st.session_state.conversation_id = None
         st.session_state.messages = []
 
     genie_space_id = st.text_input(
-        "Genie Space ID", placeholder="01efe16a65e21836acefb797ae6a8fe4", help="Room ID in the Genie Space URL"
+        "Genie Space ID",
+        placeholder="01efe16a65e21836acefb797ae6a8fe4",
+        help="Room ID in the Genie Space URL",
     )
     if genie_space_id != st.session_state.get("genie_space_id", ""):
         reset_conversation()
         st.session_state.genie_space_id = genie_space_id
 
-
     def display_message(message: Dict):
         if "content" in message:
             st.markdown(message["content"])
@@ -42,19 +43,21 @@ def display_message(message: Dict):
             with st.expander("Show generated code"):
                 st.code(message["code"], language="sql", wrap_lines=True)
 
-
-    def get_query_result(statement_id: str) -> pd.DataFrame:     
+    def get_query_result(statement_id: str) -> pd.DataFrame:
         query = w.statement_execution.get_statement(statement_id)
         result = query.result.data_array
 
         next_chunk = query.result.next_chunk_index
         while next_chunk:
-            chunk = w.statement_execution.get_statement_result_chunk_n(statement_id, next_chunk)
+            chunk = w.statement_execution.get_statement_result_chunk_n(
+                statement_id, next_chunk
+            )
             result.append(chunk.data_array)
             next_chunk = chunk.next_chunk_index
 
-        return pd.DataFrame(result, columns=[i.name for i in query.manifest.schema.columns])
-
+        return pd.DataFrame(
+            result, columns=[i.name for i in query.manifest.schema.columns]
+        )
 
     def process_genie_response(response: GenieMessage):
         st.session_state.conversation_id = response.conversation_id
@@ -67,12 +70,14 @@ def process_genie_response(response: GenieMessage):
             elif i.query:
                 data = get_query_result(i.query.statement_id)
                 message = {
-                    "role": "assistant", "content": i.query.description, "data": data, "code": i.query.query
+                    "role": "assistant",
+                    "content": i.query.description,
+                    "data": data,
+                    "code": i.query.query,
                 }
                 display_message(message)
                 st.session_state.messages.append(message)
 
-
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
@@ -93,25 +98,36 @@ def process_genie_response(response: GenieMessage):
                             genie_space_id, st.session_state.conversation_id, prompt
                         )
                     except Exception as e:
-                        status.update(label="Conversation failed. Check the required permissions.", state="error")
+                        status.update(
+                            label="Conversation failed. Check the required permissions.",
+                            state="error",
+                        )
                         st.button("New Chat", on_click=reset_conversation)
                         raise e
                     if conversation.error:
                         st.error(conversation.error.type, conversation.error.error)
                     process_genie_response(conversation)
                 else:
                     try:
-                        conversation = w.genie.start_conversation_and_wait(genie_space_id, prompt)
+                        conversation = w.genie.start_conversation_and_wait(
+                            genie_space_id, prompt
+                        )
                     except Exception as e:
-                        status.update(label="Failed to initialize Genie. Check the required permissions.", state="error")
+                        status.update(
+                            label="Failed to initialize Genie. Check the required permissions.",
+                            state="error",
+                        )
                         st.button("New Chat", on_click=reset_conversation)
                         raise e
                     if conversation.error:
                         st.error(conversation.error.type, conversation.error.error)
                     process_genie_response(conversation)
                 status.update(label="", state="complete")
                 st.button("New Chat", on_click=reset_conversation)
-                st.link_button("Open Genie", f"{w.config.host}/genie/rooms/{genie_space_id}/chats/{st.session_state.conversation_id}")
+                st.link_button(
+                    "Open Genie",
+                    f"{w.config.host}/genie/rooms/{genie_space_id}/chats/{st.session_state.conversation_id}",
+                )
             else:
                 st.error("Please fill in the Genie Space ID.")
 
@@ -154,7 +170,7 @@ def process_genie_response(response):
             message = {"role": "assistant", "content": i.text.content}
             display_message(message)
         elif i.query:
-            data = get_query_result(response.query_result.statement_id)
+            data = get_query_result(i.query.statement_id)
             message = {
                 "role": "assistant", "content": i.query.description, "data": data, "code": i.query.query
             }
@@ -188,15 +204,15 @@ def process_genie_response(response):
             * `CAN VIEW` the Genie Space
             """
         )
-    
+
     with col2:
         st.markdown(
             """
             **Databricks resources**
             * Genie API
             """
         )
-    
+
     with col3:
         st.markdown(
             """