Merge pull request #41 from antonbricks/main

Fix Genie API SDK for fetching statement_id

Author: antonbricks

dash/pages/genie_api.py

@@ -44,7 +44,7 @@ def process_genie_response(response):
         if i.text:
             print(f"A: {i.text.content}")
         elif i.query:
-            data = get_query_result(i.query.statement_id)
+            data = get_query_result(response.query_result.statement_id)
             print(f"A: {i.query.description}")
             print(f"Data: {data}")
             print(f"Generated code: {i.query.query}")
@@ -146,7 +146,7 @@ def process_genie_response(response: GenieMessage, chat_history: List[Dict]) ->
                 "content": attachment.text.content
             })
         elif attachment.query:
-            data_table = get_query_result(attachment.query.statement_id)
+            data_table = get_query_result(response.query_result.statement_id)
             chat_history.append({
                 "role": "assistant",
                 "content": f"Query: {attachment.query.description}\nGenerated SQL: {attachment.query.query}",

docs/docs/dash/bi/genie_api.mdx

@@ -29,7 +29,7 @@ def process_genie_response(response):
         if i.text:
             print(f"A: {i.text.content}")
         elif i.query:
-            data = get_query_result(i.query.statement_id)
+            data = get_query_result(response.query_result.statement_id)
             print(f"A: {i.query.description}")
             print(f"Data: {data}")
             print(f"Generated code: {i.query.query}")

docs/docs/streamlit/bi/genie_api.mdx

@@ -45,7 +45,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/genie_api.py

@@ -68,7 +68,7 @@ def process_genie_response(response: GenieMessage):
                 display_message(message)
                 st.session_state.messages.append(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,
@@ -170,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(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
             }