feat(wren-ai-service): flexible chart generation (ai-env-changed)

deployment/kustomizations/base/cm.yaml

@@ -196,6 +196,10 @@ data:
         document_store: qdrant
       - name: sql_tables_extraction
         llm: litellm_llm.default
+      - name: chart_validation
+        llm: litellm_llm.default
+      - name: chart_correction
+        llm: litellm_llm.default
 
     ---
     settings:

docker/config.example.yaml

@@ -146,6 +146,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.anthropic.yaml

@@ -140,6 +140,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.azure.yaml

@@ -153,6 +153,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.deepseek.yaml

@@ -163,6 +163,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.google_ai_studio.yaml

@@ -149,6 +149,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.grok.yaml

@@ -145,6 +145,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.groq.yaml

@@ -144,6 +144,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.lm_studio.yaml

@@ -143,6 +143,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.ollama.yaml

@@ -143,6 +143,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/docs/config_examples/config.open_router.yaml

@@ -141,6 +141,10 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
+  - name: chart_validation
+    llm: litellm_llm.default
+  - name: chart_correction
+    llm: litellm_llm.default
 
 ---
 settings:

wren-ai-service/pyproject.toml

@@ -36,6 +36,7 @@ jsonschema = "^4.23.0"
 litellm = "^1.70.0"
 boto3 = "^1.35.90"
 qdrant-client = "==1.11.0"
+vl-convert-python = "^1.7.0"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.7.1"

wren-ai-service/src/config.py

@@ -42,6 +42,8 @@ class Settings(BaseSettings):
     allow_sql_functions_retrieval: bool = Field(default=True)
     max_histories: int = Field(default=5)
     max_sql_correction_retries: int = Field(default=3)
+    allow_chart_validation: bool = Field(default=False)
+    max_chart_correction_retries: int = Field(default=3)
 
     # engine config
     engine_timeout: float = Field(default=30.0)

wren-ai-service/src/globals.py

@@ -160,7 +160,15 @@ def create_service_container(
                 "chart_generation": generation.ChartGeneration(
                     **pipe_components["chart_generation"],
                 ),
+                "chart_validation": generation.ChartValidation(
+                    **pipe_components["chart_validation"],
+                ),
+                "chart_correction": generation.ChartCorrection(
+                    **pipe_components["chart_correction"],
+                ),
             },
+            allow_chart_validation=settings.allow_chart_validation,
+            max_chart_correction_retries=settings.max_chart_correction_retries,
             **query_cache,
         ),
         chart_adjustment_service=services.ChartAdjustmentService(
@@ -172,7 +180,15 @@ def create_service_container(
                 "chart_adjustment": generation.ChartAdjustment(
                     **pipe_components["chart_adjustment"],
                 ),
+                "chart_validation": generation.ChartValidation(
+                    **pipe_components["chart_validation"],
+                ),
+                "chart_correction": generation.ChartCorrection(
+                    **pipe_components["chart_correction"],
+                ),
             },
+            allow_chart_validation=settings.allow_chart_validation,
+            max_chart_correction_retries=settings.max_chart_correction_retries,
             **query_cache,
         ),
         sql_answer_service=services.SqlAnswerService(

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

@@ -1,5 +1,7 @@
 from .chart_adjustment import ChartAdjustment
+from .chart_correction import ChartCorrection
 from .chart_generation import ChartGeneration
+from .chart_validation import ChartValidation
 from .data_assistance import DataAssistance
 from .followup_sql_generation import FollowUpSQLGeneration
 from .followup_sql_generation_reasoning import FollowUpSQLGenerationReasoning
@@ -36,4 +38,6 @@
     "FollowUpSQLGenerationReasoning",
     "MisleadingAssistance",
     "SQLTablesExtraction",
+    "ChartValidation",
+    "ChartCorrection",
 ]