You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@superset.apache.org by "SkinnyPigeon (via GitHub)" <gi...@apache.org> on 2025/11/19 13:42:54 UTC
[PR] feat(datasets): ai assistant for queries [superset]
SkinnyPigeon opened a new pull request, #36187:
URL: https://github.com/apache/superset/pull/36187
<!---
Please write the PR title following the conventions at https://www.conventionalcommits.org/en/v1.0.0/
Example:
fix(dashboard): load charts correctly
-->
### SUMMARY
AI-powered SQL generation POC for Apache Superset. Adds an AI Assistant that generates and executes SQL queries from natural language using HuggingFace models. Only the input text and the dataset's column descriptions are used to generate the query with the AI not able to view the underlying data.
I wanted to gauge interest in a feature like this in the OSS version. I know we have team members hoping to use something like this to interact with data.
---
This adds a new AI Assistant page to the Superset frontend alongside a new API endpoint. This uses the `huggingface_hub` package alongside a free token to generate the query text based on user questions about individual datasets. This works well when teams have added column descriptions, as they are passed to the model to inform it of the best way to generate the query. The videos below show it running as well as giving a quick preview of the context shared with the model.
I am not a frontend person so please forgive the sytling. If this was a feature others were interested in, I would require help getting this up to the project's standards, including help with the frontend testing suite.
I have tested this with PostgreSQL and Redshift-based datasets. More testing with other dialects would be welcome.
As the default model, I've gone for `Qwen/Qwen2.5-Coder-32B-Instruct` which seems to do the job quite well. This and the token are able to be set via env vars and have been added to the config.
### BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
Video showing the user using natural language in a chat box, which returns a SQL query that is then executed against the data source with the results shown inside the chat box
https://github.com/user-attachments/assets/aa098f92-4637-41bc-89b8-5ff2578fa7c1
Data showing the query context that allows the SQL to be generated. This is autopopulated via the `desscriptions` of the columns of the Dataset, along with additional info such as the database's dialect to try to generate the correct syntax
https://github.com/user-attachments/assets/728b53df-6f24-43bf-9c80-0970cb308ca7
The default page look:
<img width="1506" height="842" alt="Screenshot 2025-11-19 at 14 20 42" src="https://github.com/user-attachments/assets/adb30453-4b7f-4ef3-ac39-244a67a57bcf" />
The page with a Dataset selected from the drop down, giving the user a small preview of the data they will be querying:
<img width="1510" height="849" alt="Screenshot 2025-11-19 at 14 17 36" src="https://github.com/user-attachments/assets/df35e3e9-a2cf-4016-8622-dd71334f3338" />
### TESTING INSTRUCTIONS
1. Set HuggingFace API token via [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens). This can be a free account with no credit card required.
2. Make the env var available to your deployment process, e.g. `export HF_API_TOKEN=your_token`
3. Optional: Set model: `export HF_MODEL=Qwen/Qwen2.5-Coder-32B-Instruct`
4. Build and run Superset
5. Pick a dataset and add descriptions for any columns you would like to query, an example of values for `ab_user` can be found below
6. Navigate to `/aiassistant/`
7. Select a dataset
8. Ask: "Who is the most active user currently?" or "Which was the month with the most new users?"...
9. Click "Execute Query" to see results
### ADDITIONAL INFORMATION
<!--- Check any relevant boxes with "x" -->
<!--- HINT: Include "Fixes #nnn" if you are fixing an existing issue -->
- [ ] Has associated issue:
- [ ] Required feature flags:
- [x] Changes UI
- [ ] Includes DB Migration (follow approval process in [SIP-59](https://github.com/apache/superset/issues/13351))
- [ ] Migration is atomic, supports rollback & is backwards-compatible
- [ ] Confirm DB migration upgrade and downgrade tested
- [ ] Runtime estimates and downtime expectations provided
- [x] Introduces new feature or API
- [ ] Removes existing feature or API
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(datasets): ai assistant for queries [superset]
Posted by "codecov[bot] (via GitHub)" <gi...@apache.org>.
codecov[bot] commented on PR #36187:
URL: https://github.com/apache/superset/pull/36187#issuecomment-3552818581
## [Codecov](https://app.codecov.io/gh/apache/superset/pull/36187?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `50.00000%` with `76 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 48.21%. Comparing base ([`53b9045`](https://app.codecov.io/gh/apache/superset/commit/53b90459437ae59730440335a734accffb04eaa9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`0810d24`](https://app.codecov.io/gh/apache/superset/commit/0810d24d7cfc3ab696fccffc59f3d731b091bd64?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 2 commits behind head on master.
| [Files with missing lines](https://app.codecov.io/gh/apache/superset/pull/36187?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [superset/aiassistant/sql\_generator.py](https://app.codecov.io/gh/apache/superset/pull/36187?src=pr&el=tree&filepath=superset%2Faiassistant%2Fsql_generator.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvYWlhc3Npc3RhbnQvc3FsX2dlbmVyYXRvci5weQ==) | 25.00% | [33 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/36187?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [superset/aiassistant/api.py](https://app.codecov.io/gh/apache/superset/pull/36187?src=pr&el=tree&filepath=superset%2Faiassistant%2Fapi.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvYWlhc3Npc3RhbnQvYXBpLnB5) | 47.16% | [28 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/36187?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [superset/views/aiassistant.py](https://app.codecov.io/gh/apache/superset/pull/36187?src=pr&el=tree&filepath=superset%2Fviews%2Faiassistant.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvdmlld3MvYWlhc3Npc3RhbnQucHk=) | 63.41% | [15 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/36187?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #36187 +/- ##
===========================================
+ Coverage 0 48.21% +48.21%
===========================================
Files 0 633 +633
Lines 0 46355 +46355
Branches 0 5012 +5012
===========================================
+ Hits 0 22349 +22349
- Misses 0 23375 +23375
- Partials 0 631 +631
```
| [Flag](https://app.codecov.io/gh/apache/superset/pull/36187/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [hive](https://app.codecov.io/gh/apache/superset/pull/36187/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `43.91% <50.00%> (?)` | |
| [presto](https://app.codecov.io/gh/apache/superset/pull/36187/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `47.52% <50.00%> (?)` | |
| [python](https://app.codecov.io/gh/apache/superset/pull/36187/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `48.20% <50.00%> (?)` | |
| [unit](https://app.codecov.io/gh/apache/superset/pull/36187/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `100.00% <ø> (?)` | |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/superset/pull/36187?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(datasets): ai assistant for queries [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #36187:
URL: https://github.com/apache/superset/pull/36187#issuecomment-3552782654
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(datasets): ai assistant for queries [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #36187:
URL: https://github.com/apache/superset/pull/36187#discussion_r2542115030
##########
superset/aiassistant/sql_generator.py:
##########
@@ -0,0 +1,154 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""AI-powered SQL generation service using HuggingFace models."""
+
+import logging
+from typing import Any
+
+from flask import current_app
+
+logger = logging.getLogger(__name__)
+
+
+class SQLGeneratorService:
+ """Service for generating SQL from natural language using AI models."""
+
+ @staticmethod
+ def get_api_token() -> str | None:
+ """
+ Get HuggingFace API token from Superset config.
+
+ Returns:
+ API token string or None if not configured
+ """
+ return current_app.config.get("HF_API_TOKEN") or current_app.config.get(
+ "HF_TOKEN"
+ )
+
+ @staticmethod
+ def get_model() -> str:
+ """
+ Get HuggingFace model from Superset config.
+
+ Returns:
+ Model name string, defaults to Qwen/Qwen2.5-Coder-32B-Instruct
+ """
+ return current_app.config.get("HF_MODEL") or "Qwen/Qwen2.5-Coder-32B-Instruct"
+
+ @staticmethod
+ def generate_sql(
+ user_query: str, schema_info: str, db_dialect: str = "SQL"
+ ) -> dict[str, Any]:
+ """
+ Generate SQL from natural language query.
+
+ Args:
+ user_query: Natural language question from the user
+ schema_info: Schema information about the dataset
+ db_dialect: Database dialect (e.g., "PostgreSQL", "MySQL", "SQLite")
+
+ Returns:
+ Dictionary with 'sql' and optional 'error' keys
+ """
+ try:
+ from huggingface_hub import InferenceClient
+ except ImportError:
+ logger.error("huggingface_hub package is not installed")
+ return {
+ "error": (
+ "AI SQL generation is not available. "
+ "Missing huggingface_hub package."
+ )
+ }
+
+ api_token = SQLGeneratorService.get_api_token()
+
+ if not api_token:
+ logger.warning(
+ "No HuggingFace API token configured. "
+ "Set HF_API_TOKEN in superset_config.py"
+ )
+ return {
+ "error": (
+ "AI SQL generation is not configured. "
+ "Please set HF_API_TOKEN in superset_config.py"
+ )
+ }
+
+ try:
+ # Initialize the inference client
+ client = InferenceClient(token=api_token)
+
+ # Construct the messages for chat completion
+ system_prompt = (
+ f"{schema_info}\n\nIMPORTANT: Generate {db_dialect} dialect SQL. "
+ f"Use {db_dialect}-specific syntax and functions."
+ )
+ messages = [
+ {"role": "system", "content": system_prompt},
+ {
+ "role": "user",
+ "content": (
+ f"Convert this request into {db_dialect} SQL: {user_query}"
+ ),
+ },
+ ]
+
+ logger.info(
+ "Generating SQL for query: %s...", user_query[:100]
+ ) # Log first 100 chars
+
+ # Get the model to use
+ model = SQLGeneratorService.get_model()
+ logger.info("Using HuggingFace model: %s", model)
+
+ # Call the model using chat completion
+ response = client.chat_completion(
+ messages=messages,
+ model=model,
+ max_tokens=500,
+ temperature=0.1, # Low temperature for deterministic output
+ )
+
+ # Extract the SQL query
+ sql_query = response.choices[0].message.content.strip()
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Missing response validation in AI API call</b></div>
<div id="fix">
The code directly accesses `response.choices[0].message.content` without verifying that the `choices` list exists and contains at least one element. If the HuggingFace API returns a response without choices, this will raise an `IndexError`, causing the SQL generation to fail unexpectedly. In the `generate_sql` method, the API endpoints in `superset/aiassistant/api.py` and `superset/views/aiassistant.py` that call this function will receive an unhandled exception instead of a proper error response. To fix this, add a check after the API call to ensure `response.choices` is not empty before proceeding.
</div>
<details>
<summary>
<b>Code suggestion</b>
</summary>
<blockquote>Check the AI-generated fix before applying</blockquote>
<div id="code">
````suggestion
)
# Validate response structure
if not response.choices or len(response.choices) == 0:
logger.error("Invalid response from AI model: no choices returned")
return {"error": "AI model returned invalid response"}
# Extract the SQL query
sql_query = response.choices[0].message.content.strip()
````
</div>
</details>
</div>
<small><i>Code Review Run <a href=https://github.com/apache/superset/pull/36187#issuecomment-3552782654>#565ab4</a></i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
##########
superset/views/aiassistant.py:
##########
@@ -0,0 +1,90 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import logging
+from typing import Any
+
+from flask import request
+from flask_appbuilder import permission_name
+from flask_appbuilder.api import expose
+from flask_appbuilder.security.decorators import has_access
+
+from superset import event_logger
+from superset.aiassistant.sql_generator import SQLGeneratorService
+from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP
+from superset.superset_typing import FlaskResponse
+
+from .base import BaseSupersetView
+
+logger = logging.getLogger(__name__)
+
+
+class AIAssistantView(BaseSupersetView):
+ """View for the AI Assistant page."""
+
+ route_base = "/aiassistant"
+ class_permission_name = "AIAssistant"
+
+ method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP
+
+ @expose("/", methods=["GET"])
+ @has_access
+ @permission_name("read")
+ @event_logger.log_this
+ def root(self, **kwargs: Any) -> FlaskResponse:
+ """Handles the AI Assistant page."""
+ return self.render_app_template()
+
+ @expose("/generate_sql", methods=["POST"])
+ @has_access
+ @permission_name("read")
+ @event_logger.log_this
+ def generate_sql(self) -> FlaskResponse:
+ """
+ Generate SQL from natural language query using AI.
+
+ Expected JSON payload:
+ {
+ "user_query": "Show me all active users",
+ "schema_info": "The dataset 'users' has columns: id (INTEGER), ..."
+ }
+
+ Returns:
+ JSON response with 'sql' or 'error' key
+ """
+ try:
+ data = request.get_json(force=True)
+ user_query = data.get("user_query")
+ schema_info = data.get("schema_info")
+
+ if not user_query or not schema_info:
+ return self.json_response(
+ {"error": "Missing required fields: user_query and schema_info"},
+ status=400,
+ )
+
+ # Generate SQL using the AI service
+ result = SQLGeneratorService.generate_sql(user_query, schema_info)
+
+ if "error" in result:
+ logger.warning("AI SQL generation failed")
+ return self.json_response(result, status=400)
+
+ return self.json_response(result)
+
+ except Exception as e:
+ logger.exception("Error in AI SQL generation endpoint")
+ return self.json_response({"error": str(e)}, status=500)
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Insecure exception handling in generate_sql</b></div>
<div id="fix">
The exception handler in the `generate_sql` method returns the full exception message `str(e)` to the client, which could expose sensitive internal details like database connection strings or stack traces in production. While the `logger.exception` call properly logs the full error for debugging, the client response should use a generic error message to prevent information disclosure. This affects the `generate_sql` endpoint and any downstream consumers of the AI SQL generation feature.
</div>
<details>
<summary>
<b>Code suggestion</b>
</summary>
<blockquote>Check the AI-generated fix before applying</blockquote>
<div id="code">
````suggestion
except Exception as e:
logger.exception("Error in AI SQL generation endpoint")
return self.json_response({"error": "Internal server error"}, status=500)
````
</div>
</details>
</div>
<small><i>Code Review Run <a href=https://github.com/apache/superset/pull/36187#issuecomment-3552782654>#565ab4</a></i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
##########
superset-frontend/src/pages/AIAssistant/index.tsx:
##########
@@ -0,0 +1,714 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { useEffect, useState } from 'react';
+import { SupersetClient, t } from '@superset-ui/core';
+import { styled, useTheme } from '@apache-superset/core/ui';
+import {
+ Row,
+ Col,
+ Card,
+ Input,
+ Select,
+ Button,
+} from '@superset-ui/core/components';
+
+const { TextArea } = Input;
+
+const PageContainer = styled.div`
+ padding: 48px;
+ max-width: 1400px;
+ margin: 0 auto;
+`;
+
+const PageTitle = styled.h1`
+ font-size: 32px;
+ font-weight: 600;
+ margin-bottom: 48px;
+ display: flex;
+ align-items: center;
+ gap: 16px;
+`;
+
+const StyledCard = styled(Card)`
+ margin-bottom: 32px;
+ box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08);
+`;
+
+const ChatContainer = styled.div`
+ ${({ theme }) => `
+ min-height: 400px;
+ max-height: 400px;
+ overflow-y: auto;
+ padding: 32px;
+ border: 1px solid ${theme.colorBorder};
+ border-radius: 8px;
+ background-color: ${theme.colorBgLayout};
+ margin-bottom: 32px;
+ `}
+`;
+
+const ChatMessage = styled.div<{ isUser?: boolean }>`
+ ${({ theme, isUser }) => `
+ padding: 24px;
+ margin-bottom: 16px;
+ border-radius: 8px;
+ background-color: ${isUser ? theme.colorInfoBg : theme.colorBgContainer};
+ max-width: 80%;
+ ${isUser ? 'margin-left: auto;' : 'margin-right: auto;'}
+ white-space: pre-wrap;
+ word-wrap: break-word;
+
+ pre {
+ background-color: ${theme.colorBgLayout};
+ padding: 12px;
+ border-radius: 4px;
+ overflow-x: auto;
+ margin: 8px 0;
+ }
+
+ code {
+ font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
+ font-size: 13px;
+ }
+ `}
+`;
+
+const LoadingMessage = styled.div`
+ ${({ theme }) => `
+ padding: 24px;
+ margin-bottom: 16px;
+ border-radius: 8px;
+ background-color: ${theme.colorBgContainer};
+ max-width: 80%;
+ margin-right: auto;
+ color: ${theme.colorTextSecondary};
+ font-style: italic;
+ `}
+`;
+
+const DatasetSelectorContainer = styled.div`
+ margin-bottom: 32px;
+`;
+
+const StyledTextArea = styled(TextArea)`
+ margin-bottom: 16px;
+`;
+
+const DataPreviewContainer = styled.div`
+ ${({ theme }) => `
+ max-height: 500px;
+ overflow: auto;
+ border: 1px solid ${theme.colorBorder};
+ border-radius: 8px;
+ margin-top: 16px;
+ `}
+`;
+
+const ColumnTable = styled.table`
+ ${({ theme }) => `
+ width: 100%;
+ border-collapse: collapse;
+
+ th,
+ td {
+ padding: 12px;
+ text-align: left;
+ border-bottom: 1px solid ${theme.colorBorder};
+ white-space: nowrap;
+ }
+
+ th {
+ font-weight: 600;
+ background-color: ${theme.colorFillQuaternary};
+ position: sticky;
+ top: 0;
+ z-index: 1;
+ }
+
+ tr:hover {
+ background-color: ${theme.colorBgLayout};
+ }
+ `}
+`;
+
+const SQLCodeBlock = styled.pre`
+ ${({ theme }) => `
+ background-color: ${theme.colorBgLayout};
+ padding: 16px;
+ border-radius: 4px;
+ overflow-x: auto;
+ margin: 8px 0;
+ border: 1px solid ${theme.colorBorder};
+
+ code {
+ font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
+ font-size: 13px;
+ line-height: 1.5;
+ color: ${theme.colorText};
+ }
+ `}
+`;
+
+interface Dataset {
+ id: number;
+ table_name: string;
+ database_name: string;
+ schema?: string;
+}
+
+interface ChatMessageType {
+ id: string;
+ text: string;
+ isUser: boolean;
+ timestamp: Date;
+ sql?: string; // Optional SQL code to display
+ queryResults?: {
+ columns: string[];
+ data: Array<Record<string, any>>;
+ rowCount: number;
+ };
+ isExecuting?: boolean;
+}
+
+interface DatasetData {
+ columns: Array<{ name: string; type: string; description?: string }>;
+ data: Array<Record<string, any>>;
+}
+
+export default function AIAssistant() {
+ const theme = useTheme();
+ const [datasets, setDatasets] = useState<Dataset[]>([]);
+ const [selectedDataset, setSelectedDataset] = useState<number | null>(null);
+ const [datasetData, setDatasetData] = useState<DatasetData | null>(null);
+ const [loadingDatasets, setLoadingDatasets] = useState(false);
+ const [loadingData, setLoadingData] = useState(false);
+ const [loadingAI, setLoadingAI] = useState(false);
+ const [messages, setMessages] = useState<ChatMessageType[]>([]);
+ const [inputValue, setInputValue] = useState('');
+
+ // Fetch datasets on mount
+ useEffect(() => {
+ fetchDatasets();
+ }, []);
+
+ // Fetch dataset data when selection changes
+ useEffect(() => {
+ if (selectedDataset) {
+ fetchDatasetData(selectedDataset);
+ } else {
+ setDatasetData(null);
+ }
+ }, [selectedDataset]);
+
+ const fetchDatasets = async () => {
+ setLoadingDatasets(true);
+ try {
+ const response = await SupersetClient.get({
+ endpoint: '/api/v1/dataset/',
+ searchParams: {
+ q: JSON.stringify({
+ page_size: 100,
+ order_column: 'table_name',
+ order_direction: 'asc',
+ }),
+ },
+ });
+
+ const datasetsData = response.json.result.map((ds: any) => ({
+ id: ds.id,
+ table_name: ds.table_name,
+ database_name: ds.database?.database_name || 'Unknown',
+ schema: ds.schema,
+ }));
+
+ setDatasets(datasetsData);
+ } catch (error) {
+ console.error('Error fetching datasets:', error);
+ } finally {
+ setLoadingDatasets(false);
+ }
+ };
+
+ const fetchDatasetData = async (datasetId: number) => {
+ setLoadingData(true);
+ try {
+ // First, fetch dataset metadata to get column information
+ const datasetResponse = await SupersetClient.get({
+ endpoint: `/api/v1/dataset/${datasetId}`,
+ });
+
+ const dataset = datasetResponse.json.result;
+
+ // Get column information including descriptions
+ const columns =
+ dataset.columns?.map((col: any) => ({
+ name: col.column_name,
+ type: col.type,
+ description: col.description,
+ })) || [];
+
+ // Log column info for debugging
+ console.log('Dataset columns with descriptions:', columns);
+
+ // Now fetch sample data using the chart data API
+ const queryPayload = {
+ datasource: {
+ id: datasetId,
+ type: 'table',
+ },
+ queries: [
+ {
+ columns: columns.map((col: { name: string }) => col.name),
+ row_limit: 50,
+ orderby: [],
+ },
+ ],
+ result_format: 'json',
+ result_type: 'full',
+ };
+
+ const dataResponse = await SupersetClient.post({
+ endpoint: '/api/v1/chart/data',
+ jsonPayload: queryPayload,
+ });
+
+ // Extract the data from the response
+ const result = dataResponse.json.result?.[0];
+ const data = result?.data || [];
+
+ setDatasetData({
+ columns,
+ data,
+ });
+ } catch (error) {
+ console.error('Error fetching dataset data:', error);
+ // Still show columns even if data fetch fails
+ const datasetResponse = await SupersetClient.get({
+ endpoint: `/api/v1/dataset/${datasetId}`,
+ });
+ const dataset = datasetResponse.json.result;
+ const columns =
+ dataset.columns?.map((col: any) => ({
+ name: col.column_name,
+ type: col.type,
+ description: col.description,
+ })) || [];
+ setDatasetData({
+ columns,
+ data: [],
+ });
+ } finally {
+ setLoadingData(false);
+ }
+ };
+
+ // Build schema information for AI context
+ const buildSchemaInfo = (): string => {
+ if (!selectedDataset || !datasetData) return '';
+
+ const datasetName =
+ datasets.find(d => d.id === selectedDataset)?.table_name || 'Unknown';
+ let schemaInfo = `The dataset "${datasetName}" has the following columns:\n\n`;
+
+ datasetData.columns.forEach(col => {
+ schemaInfo += `- ${col.name} (${col.type})`;
+ if (col.description) {
+ schemaInfo += `: ${col.description}`;
+ }
+ schemaInfo += '\n';
+ });
+
+ schemaInfo +=
+ '\nGenerate a valid SQL SELECT query using ONLY these column names.\n';
+ schemaInfo +=
+ 'DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself.';
+
+ return schemaInfo;
+ };
+
+ const handleSendMessage = async () => {
+ if (!inputValue.trim()) return;
+
+ if (!selectedDataset) {
+ // Show error if no dataset selected
+ const errorMessage: ChatMessageType = {
+ id: Date.now().toString(),
+ text: 'Please select a dataset first before asking questions.',
+ isUser: false,
+ timestamp: new Date(),
+ };
+ setMessages(prev => [...prev, errorMessage]);
+ return;
+ }
+
+ const newMessage: ChatMessageType = {
+ id: Date.now().toString(),
+ text: inputValue,
+ isUser: true,
+ timestamp: new Date(),
+ };
+
+ setMessages(prev => [...prev, newMessage]);
+ setInputValue('');
+ setLoadingAI(true);
+
+ try {
+ // Call the AI SQL generation API
+ const response = await SupersetClient.post({
+ endpoint: '/api/v1/aiassistant/generate_sql',
+ jsonPayload: {
+ dataset_id: selectedDataset,
+ user_query: inputValue,
+ },
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Incomplete API payload for SQL generation</b></div>
<div id="fix">
The API call to generate SQL is missing the required 'schema_info' parameter that the backend expects. This will cause the request to fail with a 400 error indicating missing required fields. The backend endpoint requires both 'user_query' and 'schema_info' to generate SQL, but the frontend is only sending 'dataset_id' and 'user_query'. Add the schema_info built from the selected dataset's metadata to the jsonPayload.
</div>
<details>
<summary>
<b>Code suggestion</b>
</summary>
<blockquote>Check the AI-generated fix before applying</blockquote>
<div id="code">
````suggestion
endpoint: '/api/v1/aiassistant/generate_sql',
jsonPayload: {
dataset_id: selectedDataset,
user_query: inputValue,
schema_info: buildSchemaInfo(),
},
````
</div>
</details>
</div>
<small><i>Code Review Run <a href=https://github.com/apache/superset/pull/36187#issuecomment-3552782654>#565ab4</a></i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(datasets): ai assistant for queries [superset]
Posted by "SkinnyPigeon (via GitHub)" <gi...@apache.org>.
SkinnyPigeon closed pull request #36187: feat(datasets): ai assistant for queries
URL: https://github.com/apache/superset/pull/36187
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org