fix chart adjustment

Author: cyyeh

wren-ai-service/src/globals.py

@@ -262,7 +262,7 @@ def create_service_container(
         ),
         conversation_service=v2_services.ConversationService(
             pipelines={
-                "intent_classification": generation.IntentClassification(
+                "intent_classification": generation.IntentClassificationV2(
                     **pipe_components["intent_classification"],
                     wren_ai_docs=wren_ai_docs,
                 ),
@@ -339,7 +339,7 @@ def create_service_container(
                 "chart_generation": generation.ChartGeneration(
                     **pipe_components["chart_generation"],
                 ),
-                "chart_adjustment": generation.ChartAdjustment(
+                "chart_adjustment": generation.ChartAdjustmentV2(
                     **pipe_components["chart_adjustment"],
                 ),
             },

wren-ai-service/src/pipelines/generation/__init__.py

@@ -1,10 +1,12 @@
 from .chart_adjustment import ChartAdjustment
+from .chart_adjustment_v2 import ChartAdjustmentV2
 from .chart_generation import ChartGeneration
 from .data_assistance import DataAssistance
 from .data_exploration_assistance import DataExplorationAssistance
 from .followup_sql_generation import FollowUpSQLGeneration
 from .followup_sql_generation_reasoning import FollowUpSQLGenerationReasoning
 from .intent_classification import IntentClassification
+from .intent_classification_v2 import IntentClassificationV2
 from .misleading_assistance import MisleadingAssistance
 from .question_recommendation import QuestionRecommendation
 from .relationship_recommendation import RelationshipRecommendation
@@ -20,9 +22,11 @@
 __all__ = [
     "ChartGeneration",
     "ChartAdjustment",
+    "ChartAdjustmentV2",
     "DataAssistance",
     "FollowUpSQLGeneration",
     "IntentClassification",
+    "IntentClassificationV2",
     "QuestionRecommendation",
     "RelationshipRecommendation",
     "SemanticsDescription",

wren-ai-service/src/pipelines/generation/chart_adjustment_v2.py

@@ -0,0 +1,182 @@
+import logging
+import sys
+from typing import Any, Dict
+
+import orjson
+from hamilton import base
+from hamilton.async_driver import AsyncDriver
+from haystack.components.builders.prompt_builder import PromptBuilder
+from langfuse.decorators import observe
+
+from src.core.pipeline import BasicPipeline
+from src.core.provider import LLMProvider
+from src.pipelines.generation.utils.chart import (
+    ChartDataPreprocessor,
+    ChartGenerationPostProcessor,
+    ChartGenerationResults,
+    chart_generation_instructions,
+)
+
+logger = logging.getLogger("wren-ai-service")
+
+
+chart_adjustment_system_prompt = f"""
+### TASK ###
+
+You are a data analyst great at visualizing data using vega-lite! Given the user's request, SQL, sample data, sample column values, original vega-lite schema, 
+you need to re-generate vega-lite schema in JSON and provide suitable chart type.
+Besides, you need to give a concise and easy-to-understand reasoning to describe why you provide such vega-lite schema based on the question, SQL, sample data, sample column values, original vega-lite schema and adjustment options.
+
+{chart_generation_instructions}
+- If you think the user's request are not suitable for the data or not suitable for generating the chart, you can return an empty string for the schema and chart type and give reasoning to explain why.
+
+### OUTPUT FORMAT ###
+
+Please provide your chain of thought reasoning, chart type and the vega-lite schema in JSON format.
+
+{{
+    "reasoning": <REASON_TO_CHOOSE_THE_SCHEMA_IN_STRING_FORMATTED_IN_LANGUAGE_PROVIDED_BY_USER>,
+    "chart_type": "line" | "multi_line" | "bar" | "pie" | "grouped_bar" | "stacked_bar" | "area" | "",
+    "chart_schema": <VEGA_LITE_JSON_SCHEMA>
+}}
+"""
+
+chart_adjustment_user_prompt_template = """
+### INPUT ###
+User's request: {{ query }}
+User's SQL: {{ sql }}
+User's Vega-Lite Schema: {{ chart_schema }}
+Sample Data: {{ sample_data }}
+Sample Column Values: {{ sample_column_values }}
+Language: {{ language }}
+
+Please think step by step
+"""
+
+
+## Start of Pipeline
+@observe(capture_input=False)
+def preprocess_data(
+    data: Dict[str, Any], chart_data_preprocessor: ChartDataPreprocessor
+) -> dict:
+    return chart_data_preprocessor.run(data)
+
+
+@observe(capture_input=False)
+def prompt(
+    query: str,
+    sql: str,
+    chart_schema: dict,
+    preprocess_data: dict,
+    language: str,
+    prompt_builder: PromptBuilder,
+) -> dict:
+    sample_data = preprocess_data.get("sample_data")
+    sample_column_values = preprocess_data.get("sample_column_values")
+
+    return prompt_builder.run(
+        query=query,
+        sql=sql,
+        chart_schema=chart_schema,
+        sample_data=sample_data,
+        sample_column_values=sample_column_values,
+        language=language,
+    )
+
+
+@observe(as_type="generation", capture_input=False)
+async def generate_chart_adjustment(prompt: dict, generator: Any) -> dict:
+    return await generator(prompt=prompt.get("prompt"))
+
+
+@observe(capture_input=False)
+def post_process(
+    generate_chart_adjustment: dict,
+    vega_schema: Dict[str, Any],
+    preprocess_data: dict,
+    post_processor: ChartGenerationPostProcessor,
+) -> dict:
+    return post_processor.run(
+        generate_chart_adjustment.get("replies"),
+        vega_schema,
+        preprocess_data["sample_data"],
+    )
+
+
+## End of Pipeline
+CHART_ADJUSTMENT_MODEL_KWARGS = {
+    "response_format": {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "chart_adjustment_results",
+            "schema": ChartGenerationResults.model_json_schema(),
+        },
+    }
+}
+
+
+class ChartAdjustmentV2(BasicPipeline):
+    def __init__(
+        self,
+        llm_provider: LLMProvider,
+        **kwargs,
+    ):
+        self._components = {
+            "prompt_builder": PromptBuilder(
+                template=chart_adjustment_user_prompt_template
+            ),
+            "generator": llm_provider.get_generator(
+                system_prompt=chart_adjustment_system_prompt,
+                generation_kwargs=CHART_ADJUSTMENT_MODEL_KWARGS,
+            ),
+            "chart_data_preprocessor": ChartDataPreprocessor(),
+            "post_processor": ChartGenerationPostProcessor(),
+        }
+
+        with open("src/pipelines/generation/utils/vega-lite-schema-v5.json", "r") as f:
+            _vega_schema = orjson.loads(f.read())
+
+        self._configs = {
+            "vega_schema": _vega_schema,
+        }
+        super().__init__(
+            AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
+        )
+
+    @observe(name="Chart Adjustment")
+    async def run(
+        self,
+        query: str,
+        sql: str,
+        chart_schema: dict,
+        data: dict,
+        language: str,
+    ) -> dict:
+        logger.info("Chart Adjustment pipeline is running...")
+
+        return await self._pipe.execute(
+            ["post_process"],
+            inputs={
+                "query": query,
+                "sql": sql,
+                "chart_schema": chart_schema,
+                "data": data,
+                "language": language,
+                **self._components,
+                **self._configs,
+            },
+        )
+
+
+if __name__ == "__main__":
+    from src.pipelines.common import dry_run_pipeline
+
+    dry_run_pipeline(
+        ChartAdjustmentV2,
+        "chart_adjustment",
+        query="show me the dataset",
+        sql="",
+        chart_schema={},
+        data={},
+        language="English",
+    )

wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py

@@ -63,9 +63,9 @@
 
 ### User's QUERY HISTORY ###
 {% for history in histories %}
-User's Question:
+Question:
 {{ history.question }}
-Assistant's Response:
+SQL:
 {{ history.sql }}
 {% endfor %}
 

wren-ai-service/src/pipelines/generation/intent_classification.py

@@ -23,8 +23,7 @@
 
 intent_classification_system_prompt = """
 ### Task ###
-You are an expert detective specializing in intent classification. Combine the user's current question and previous questions to determine their true intent based on the provided database schema or sql data if provided.
-Classify the intent into one of these categories: `MISLEADING_QUERY`, `TEXT_TO_SQL`, `DATA_EXPLORATION`, `GENERAL`, `CHART`, or `USER_GUIDE`. Additionally, provide a concise reasoning (maximum 20 words) for your classification.
+You are an expert detective specializing in intent classification. Combine the user's current question and previous questions to determine their true intent based on the provided database schema. Classify the intent into one of these categories: `MISLEADING_QUERY`, `TEXT_TO_SQL`, `GENERAL`, or `USER_GUIDE`. Additionally, provide a concise reasoning (maximum 20 words) for your classification.
 
 ### Instructions ###
 - **Follow the user's previous questions:** If there are previous questions, try to understand the user's current question as following the previous questions.
@@ -36,20 +35,6 @@
 
 ### Intent Definitions ###
 
-<DATA_EXPLORATION>
-**When to Use:**
-- The user's question is about data exploration such as asking for data details, asking for explanation of the data, asking for insights, asking for recommendations, asking for comparison, etc.
-
-**Requirements:**
-- SQL DATA is provided and the user's question is about exploring the data.
-- The user's question can be answered by the SQL DATA.
-
-**Examples:**  
-- "Show me the part where the data appears abnormal"
-- "Please explain the data in the table"
-- "What's the trend of the data?"
-</DATA_EXPLORATION>
-
 <TEXT_TO_SQL>
 **When to Use:**  
 - The user's inputs are about modifying SQL from previous questions.
@@ -59,27 +44,13 @@
 **Requirements:**
 - Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions.
 - Reference phrases from the user's inputs that clearly relate to the schema.
-- The SQL DATA is not provided or SQL DATA cannot answer the user's question, and the user's question can be answered given the database schema.
 
 **Examples:**  
 - "What is the total sales for last quarter?"
 - "Show me all customers who purchased product X."
 - "List the top 10 products by revenue."
 </TEXT_TO_SQL>
 
-<CHART>
-**When to Use:**  
-- The user's question is about generating a chart.
-
-**Requirements:**  
-- The user's question can be answered by the SQL DATA.
-- SQL DATA is provided.
-- Should pick a SQL from user query histories and the picked SQL should be reflected to the SQL DATA provided.
-
-**Examples:**  
-- "Show me the bar chart of the data"
-</CHART>
-
 <GENERAL>
 **When to Use:**  
 - The user seeks general information about the database schema or its overall capabilities.
@@ -112,11 +83,9 @@
 - The user's inputs is irrelevant to the database schema or includes SQL code.
 - The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query.
 - It appears off-topic or is simply a casual conversation starter.
-- The user's question is about generating a chart but the SQL DATA is not provided.
 
 **Requirements:**  
-- For generating SQL: respond to users by incorporating phrases from the user's inputs that indicate the lack of relevance to the database schema.
-- For generating chart: respond to users that we can generate chart only if there is some data available.
+- Incorporate phrases from the user's inputs that indicate the lack of relevance to the database schema.
 
 **Examples:**  
 - "How are you?"
@@ -130,8 +99,7 @@
 {
     "rephrased_question": "<rephrased question in full standalone question if there are previous questions, otherwise the original question>",
     "reasoning": "<brief chain-of-thought reasoning (max 20 words)>",
-    "results": "MISLEADING_QUERY" | "TEXT_TO_SQL" | "GENERAL" | "USER_GUIDE" | "DATA_EXPLORATION" | "CHART",
-    "sql": "<sql query to be used for generating chart if the intent is CHART, otherwise an empty string>"
+    "results": "MISLEADING_QUERY" | "TEXT_TO_SQL" | "GENERAL" | "USER_GUIDE"
 }
 """
 
@@ -161,18 +129,13 @@
 - {{doc.path}}: {{doc.content}}
 {% endfor %}
 
-{% if sql_data %}
-### SQL DATA ###
-{{ sql_data }}
-{% endif %}
-
 ### INPUT ###
 {% if histories %}
 User's previous questions:
 {% for history in histories %}
-User's Question:
+Question:
 {{ history.question }}
-Assistant's Response:
+SQL:
 {{ history.sql }}
 {% endfor %}
 {% endif %}
@@ -296,7 +259,6 @@ def prompt(
     construct_db_schemas: list[str],
     histories: list[AskHistory],
     prompt_builder: PromptBuilder,
-    sql_data: dict,
     sql_samples: Optional[list[dict]] = None,
     instructions: Optional[list[dict]] = None,
     configuration: Configuration | None = None,
@@ -313,7 +275,6 @@ def prompt(
         ),
         current_time=configuration.show_current_time(),
         docs=wren_ai_docs,
-        sql_data=sql_data,
     )
 
 
@@ -330,15 +291,13 @@ def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict
             "rephrased_question": results["rephrased_question"],
             "intent": results["results"],
             "reasoning": results["reasoning"],
-            "sql": results["sql"],
             "db_schemas": construct_db_schemas,
         }
     except Exception:
         return {
             "rephrased_question": "",
             "intent": "TEXT_TO_SQL",
             "reasoning": "",
-            "sql": "",
             "db_schemas": construct_db_schemas,
         }
 
@@ -348,9 +307,7 @@ def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict
 
 class IntentClassificationResult(BaseModel):
     rephrased_question: str
-    results: Literal[
-        "MISLEADING_QUERY", "TEXT_TO_SQL", "GENERAL", "USER_GUIDE", "DATA_EXPLORATION"
-    ]
+    results: Literal["MISLEADING_QUERY", "TEXT_TO_SQL", "GENERAL", "USER_GUIDE"]
     reasoning: str
 
 
@@ -412,7 +369,6 @@ async def run(
         sql_samples: Optional[list[dict]] = None,
         instructions: Optional[list[dict]] = None,
         configuration: Configuration = Configuration(),
-        sql_data: Optional[dict] = None,
     ):
         logger.info("Intent Classification pipeline is running...")
         return await self._pipe.execute(
@@ -424,7 +380,6 @@ async def run(
                 "sql_samples": sql_samples or [],
                 "instructions": instructions or [],
                 "configuration": configuration,
-                "sql_data": sql_data or {},
                 **self._components,
                 **self._configs,
             },