Spaces:
Sleeping
Sleeping
| import time | |
| import requests | |
| # --- Configuration --- | |
| DATABRICKS_INSTANCE = "dbc-da15e658-afcd.cloud.databricks.com" | |
| API_TOKEN = "dapie17b5e771da1b7b481063e027e570e13" | |
| SPACE_ID = "01f1580f04271dc0aa6ad4b41298f0e4" | |
| USER_PROMPT = "What were our top 5 selling products last month by total revenue?" | |
| # Base setups | |
| BASE_URL = f"https://{DATABRICKS_INSTANCE}/api/2.0/genie/spaces/{SPACE_ID}" | |
| HEADERS = { | |
| "Authorization": f"Bearer {API_TOKEN}", | |
| "Content-Type": "application/json" | |
| } | |
| def fetch_table_data(statement_id): | |
| """Fetches the actual rows and data for the query Genie executed.""" | |
| print(f"Fetching table results for Statement: {statement_id}...") | |
| sql_url = f"https://{DATABRICKS_INSTANCE}/api/2.0/sql/statements/{statement_id}" | |
| # Poll the SQL execution warehouse if it's still finalizing the chunk | |
| while True: | |
| res = requests.get(sql_url, headers=HEADERS) | |
| res.raise_for_status() | |
| sql_data = res.json() | |
| status = sql_data.get("status", {}).get("state") | |
| if status == "SUCCEEDED": | |
| manifest = sql_data.get("manifest", {}) | |
| columns = [col["name"] for col in manifest.get("schema", {}).get("columns", [])] | |
| data_array = sql_data.get("result", {}).get("data_array", []) | |
| # Print data cleanly in a list format | |
| print(f"\nColumns: {columns}") | |
| print("-" * 50) | |
| for row in data_array: | |
| print(row) | |
| break | |
| elif status in ["FAILED", "CANCELED"]: | |
| print(f"Failed to extract structured warehouse data: {status}") | |
| break | |
| time.sleep(2) | |
| def ask_genie(): | |
| # Step 1: Start a new conversation thread inside the Space | |
| print("Initializing Genie conversation thread...") | |
| start_url = f"{BASE_URL}/start-conversation" | |
| start_payload = {"content": USER_PROMPT} | |
| response = requests.post(start_url, headers=HEADERS, json=start_payload) | |
| response.raise_for_status() | |
| convo_data = response.json() | |
| # Handle both payload response flavors | |
| conversation_id = convo_data.get("conversation_id") | |
| message_id = convo_data.get("message_id") or convo_data.get("id") | |
| print(f"Conversation Started! Convo ID: {conversation_id}") | |
| # Step 2 & 3: Poll the execution status until complete | |
| status_url = f"{BASE_URL}/conversations/{conversation_id}/messages/{message_id}" | |
| print("Waiting for Genie to process metadata and execute SQL queries...") | |
| while True: | |
| status_response = requests.get(status_url, headers=HEADERS) | |
| status_response.raise_for_status() | |
| message_status = status_response.json() | |
| current_state = message_status.get("status") | |
| print(f"Current Status: {current_state}") | |
| if current_state == "COMPLETED": | |
| print("\n--- Genie Summary Generated ---") | |
| # 1. Print Text Response Summary | |
| attachments = message_status.get("attachments", []) | |
| for attachment in attachments: | |
| if "text" in attachment: | |
| print(attachment["text"].get("content")) | |
| # 2. Extract Data Table Results using Statement ID | |
| query_metadata = message_status.get("query_result_metadata") | |
| if query_metadata and "statement_id" in query_metadata: | |
| fetch_table_data(query_metadata["statement_id"]) | |
| else: | |
| # Alternate location Databricks drops query details inside the attachments nested structures | |
| for attachment in attachments: | |
| if "query" in attachment and "query_result_metadata" in attachment["query"]: | |
| stmt_id = attachment["query"]["query_result_metadata"].get("statement_id") | |
| if stmt_id: | |
| fetch_table_data(stmt_id) | |
| break | |
| elif current_state in ["FAILED", "CANCELLED", "QUERY_RESULT_EXPIRED"]: | |
| print(f"\nExecution terminated with status: {current_state}") | |
| if "error" in message_status: | |
| print(f"Error Details: {message_status['error']}") | |
| break | |
| time.sleep(3) | |
| if __name__ == "__main__": | |
| ask_genie() |