Merge pull request #22 from ImaneTS/fix/code_snippets_genie_and_ml

Fix: Update code snippets for invoking model in streamlit and Genie room embedding

Author: pbv0

streamlit/views/genie_api.py

@@ -154,7 +154,7 @@ def process_genie_response(response):
             message = {"role": "assistant", "content": i.text.content}
             display_message(message)
         elif i.query:
-            data = get_query_result(i.query.statement_id)
+            data = get_query_result(response.query_result.statement_id)
             message = {
                 "role": "assistant", "content": i.query.description, "data": data, "code": i.query.query
             }

streamlit/views/ml_serving_invoke.py

@@ -74,16 +74,17 @@
         "code": """
         ```python
         from databricks.sdk import WorkspaceClient
+        from databricks.sdk.service.serving import DataframeSplitInput
         import streamlit as st
 
         w = WorkspaceClient()
 
         response = w.serving_endpoints.query(
             name="custom-regression-model",
-            dataframe_split={
+            dataframe_split=DataframeSplitInput.from_dict({
                 "columns": ["feature1", "feature2"],
                 "data": [[1.5, 2.5]]
-            }
+            })
         )
         st.json(response.as_dict())
         ```