Merge pull request #10 from pbv0/genie-dev-ctd

Improved Genie

Author: antonbricks

dash/pages/genie_api.py

@@ -27,6 +27,9 @@
 from databricks.sdk import WorkspaceClient
 
 
+# Refer to the source code for the full implmenetation.
+
+
 def get_query_result(statement_id):
     # For simplicity, let's say data fits in one chunk, query.manifest.total_chunk_count = 1
 
@@ -274,28 +277,31 @@ def update_chat(n_clicks, genie_space_id, conversation_id, prompt, chat_history)
     try:
         if conversation_id:
             conversation = w.genie.create_message_and_wait(genie_space_id, conversation_id, prompt)
+            # Access error codes via conversation.error for detailed exception handling.
         else:
             conversation = w.genie.start_conversation_and_wait(genie_space_id, prompt)
             conversation_id = conversation.conversation_id
+            # Access error codes via conversation.error for detailed exception handling.
 
         chat_history = process_genie_response(conversation, chat_history)
         chat_display = format_message_display(chat_history)
 
         return chat_history, chat_display, conversation_id
 
     except Exception as e:
-        return dash.no_update, dbc.Alert(f"An error occurred: {str(e)}", color="danger")
+        return dash.no_update, dbc.Alert(f"Check the required permissions. An error occurred: {str(e)}", color="danger"), ""
 
 
 @callback(
     [Output("chat-history-store", "data", allow_duplicate=True),
      Output("chat-history", "children", allow_duplicate=True),
      Output("conversation-id", "value", allow_duplicate=True),],
-     Input("clear-button", "n_clicks"),
+    [Input("clear-button", "n_clicks"),
+     Input("genie-space-id", "value"),],
      prevent_initial_call=True,
 )
-def clear_chat(n_clicks):
-    return [], [], None if n_clicks else (dash.no_update, dash.no_update, None)
+def clear_chat(n_clicks, value):
+    return [], [], None if n_clicks or value else (dash.no_update, dash.no_update, None)
 
 
 @callback(

streamlit/views/genie_api.py

@@ -1,5 +1,6 @@
 import streamlit as st
 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
@@ -20,10 +21,16 @@
 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"
     )
-    st.session_state.genie_space_id = genie_space_id
+    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):
@@ -65,18 +72,13 @@ def process_genie_response(response: GenieMessage):
                 display_message(message)
                 st.session_state.messages.append(message)
 
-    def reset_conversation():
-        st.session_state.conversation_id = None
-        st.session_state.messages = []
-
 
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
-    if st.session_state.genie_space_id == genie_space_id:
-        for message in st.session_state.messages:
-            with st.chat_message(message["role"]):
-                display_message(message)
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            display_message(message)
 
     if prompt := st.chat_input("Ask your question..."):
         st.chat_message("user").markdown(prompt)
@@ -86,22 +88,36 @@ def reset_conversation():
             if genie_space_id:
                 status = st.status("Thinking")
                 if st.session_state.get("conversation_id"):
-                    conversation = w.genie.create_message_and_wait(
-                        genie_space_id, st.session_state.conversation_id, prompt
-                    )
+                    try:
+                        conversation = w.genie.create_message_and_wait(
+                            genie_space_id, st.session_state.conversation_id, prompt
+                        )
+                    except Exception as e:
+                        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:
-                    conversation = w.genie.start_conversation_and_wait(genie_space_id, prompt)
+                    try:
+                        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")
+                        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}")
             else:
                 st.error("Please fill in the Genie Space ID.")
 
 
 with tab_b:
+    st.markdown("Refer to the source code for the full implmenetation.")
     st.code(
         """
 import streamlit as st