chore(wren-ai-service): improve text2sql

[cyyeh]

• remove unused code
• fix sql issue
• separate json prompt
• improve misleading prompt

Summary by CodeRabbit

• New Features

  • Added support for handling JSON fields in SQL generation, regeneration, and follow-up SQL pipelines, with new instructions and flags to improve query accuracy for databases containing JSON columns.

• Bug Fixes

  • Fixed the response length constraint in misleading assistance to a consistent 100 words for all languages.

• Refactor

  • Simplified and centralized JSON-related SQL instructions for better maintainability.
  • Removed unused configuration parameters and streamlined pipeline input handling.

• Chores

  • Removed standalone script execution blocks and dry-run code from multiple modules to reduce code clutter and improve maintainability.

[coderabbitai]

Walkthrough

This change refactors the handling of JSON fields in the pipeline system. It introduces a new has_json_field flag throughout schema retrieval and SQL generation pipelines, updates function and method signatures accordingly, and centralizes JSON SQL usage instructions. The update also removes the configuration parameter and all dry-run code blocks across multiple modules.

Changes

Files/Paths	Change Summary
src/core/pipeline.py	Removed async_validate function and its asyncio import.
src/pipelines/common.py	Updated build_table_ddl to return has_json_field; removed dry_run_pipeline.
src/pipelines/generation/chart_adjustment.py, .../chart_generation.py, ...	Removed all if __name__ == "__main__": dry-run/test blocks from generation pipelines.
src/pipelines/generation/data_assistance.py, .../followup_sql_generation_reasoning.py, .../intent_classification.py, .../misleading_assistance.py, .../question_recommendation.py, .../relationship_recommendation.py, .../semantics_description.py, .../sql_answer.py, .../sql_correction.py, .../sql_generation.py, .../sql_generation_reasoning.py, .../sql_question.py, .../sql_regeneration.py, .../sql_tables_extraction.py, .../user_guide_assistance.py	Removed all if __name__ == "__main__": dry-run/test blocks from generation pipelines.
src/pipelines/generation/followup_sql_generation.py, .../sql_generation.py, .../sql_regeneration.py	Added support for json_field_instructions and has_json_field flag; removed configuration parameter.
src/pipelines/generation/utils/sql.py	Moved JSON SQL usage instructions to new json_field_instructions variable; updated documentation strings.
src/pipelines/indexing/db_schema.py, .../historical_question.py, .../instructions.py, .../project_meta.py, .../sql_pairs.py, .../table_description.py	Removed all if __name__ == "__main__": dry-run/test blocks from indexing pipelines.
src/pipelines/retrieval/db_schema_retrieval.py	Added has_json_field flag tracking and propagation in schema retrieval; updated related function signatures.
src/pipelines/retrieval/historical_question_retrieval.py, .../instructions.py, .../preprocess_sql_data.py, .../sql_executor.py, .../sql_functions.py, .../sql_pairs_retrieval.py	Removed all if __name__ == "__main__": dry-run/test blocks from retrieval pipelines.
src/web/v1/services/ask.py	Extracted and passed has_json_field flag in pipeline calls; removed configuration parameter in some places.
src/web/v1/services/question_recommendation.py	Removed Configuration usage; added has_json_field handling and propagation.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant WebService
    participant RetrievalPipeline
    participant SQLGenerationPipeline

    User->>WebService: Submit query
    WebService->>RetrievalPipeline: Retrieve schema info
    RetrievalPipeline->>WebService: Return {has_calculated_field, has_metric, has_json_field}
    WebService->>SQLGenerationPipeline: Generate SQL (includes has_json_field flag)
    SQLGenerationPipeline->>WebService: Return SQL
    WebService->>User: Return results

Loading

Possibly related PRs

• Canner/WrenAI#1526: Both PRs modify SQL regeneration pipeline interfaces, with this PR removing configuration and adding has_json_field, while the related PR adds SQL functions/samples but still uses configuration.

Suggested reviewers

• yichieh-lu

Poem

In fields of JSON, rabbits hop—
New flags in code, old configs drop.
Dry-run tunnels now are sealed,
SQL rules more clearly revealed.
From schema roots to SQL leaves,
The pipeline garden softly weaves.
🐇✨

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 727f7fc and 11c642a.

📒 Files selected for processing (36)

• wren-ai-service/src/core/pipeline.py (0 hunks)
• wren-ai-service/src/pipelines/common.py (3 hunks)
• wren-ai-service/src/pipelines/generation/chart_adjustment.py (0 hunks)
• wren-ai-service/src/pipelines/generation/chart_generation.py (0 hunks)
• wren-ai-service/src/pipelines/generation/data_assistance.py (0 hunks)
• wren-ai-service/src/pipelines/generation/followup_sql_generation.py (6 hunks)
• wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py (0 hunks)
• wren-ai-service/src/pipelines/generation/intent_classification.py (1 hunks)
• wren-ai-service/src/pipelines/generation/misleading_assistance.py (1 hunks)
• wren-ai-service/src/pipelines/generation/question_recommendation.py (0 hunks)
• wren-ai-service/src/pipelines/generation/relationship_recommendation.py (0 hunks)
• wren-ai-service/src/pipelines/generation/semantics_description.py (0 hunks)
• wren-ai-service/src/pipelines/generation/sql_answer.py (0 hunks)
• wren-ai-service/src/pipelines/generation/sql_correction.py (0 hunks)
• wren-ai-service/src/pipelines/generation/sql_generation.py (6 hunks)
• wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py (0 hunks)
• wren-ai-service/src/pipelines/generation/sql_question.py (0 hunks)
• wren-ai-service/src/pipelines/generation/sql_regeneration.py (6 hunks)
• wren-ai-service/src/pipelines/generation/sql_tables_extraction.py (0 hunks)
• wren-ai-service/src/pipelines/generation/user_guide_assistance.py (0 hunks)
• wren-ai-service/src/pipelines/generation/utils/sql.py (2 hunks)
• wren-ai-service/src/pipelines/indexing/db_schema.py (0 hunks)
• wren-ai-service/src/pipelines/indexing/historical_question.py (0 hunks)
• wren-ai-service/src/pipelines/indexing/instructions.py (0 hunks)
• wren-ai-service/src/pipelines/indexing/project_meta.py (0 hunks)
• wren-ai-service/src/pipelines/indexing/sql_pairs.py (0 hunks)
• wren-ai-service/src/pipelines/indexing/table_description.py (0 hunks)
• wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (5 hunks)
• wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py (0 hunks)
• wren-ai-service/src/pipelines/retrieval/instructions.py (0 hunks)
• wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py (0 hunks)
• wren-ai-service/src/pipelines/retrieval/sql_executor.py (0 hunks)
• wren-ai-service/src/pipelines/retrieval/sql_functions.py (0 hunks)
• wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py (0 hunks)
• wren-ai-service/src/web/v1/services/ask.py (5 hunks)
• wren-ai-service/src/web/v1/services/question_recommendation.py (5 hunks)

💤 Files with no reviewable changes (26)

• wren-ai-service/src/pipelines/retrieval/sql_functions.py
• wren-ai-service/src/pipelines/generation/relationship_recommendation.py
• wren-ai-service/src/pipelines/retrieval/sql_executor.py
• wren-ai-service/src/pipelines/indexing/sql_pairs.py
• wren-ai-service/src/pipelines/generation/sql_answer.py
• wren-ai-service/src/pipelines/generation/sql_question.py
• wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py
• wren-ai-service/src/core/pipeline.py
• wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py
• wren-ai-service/src/pipelines/retrieval/instructions.py
• wren-ai-service/src/pipelines/generation/sql_correction.py
• wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py
• wren-ai-service/src/pipelines/generation/user_guide_assistance.py
• wren-ai-service/src/pipelines/indexing/table_description.py
• wren-ai-service/src/pipelines/generation/semantics_description.py
• wren-ai-service/src/pipelines/indexing/instructions.py
• wren-ai-service/src/pipelines/generation/chart_adjustment.py
• wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py
• wren-ai-service/src/pipelines/indexing/historical_question.py
• wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py
• wren-ai-service/src/pipelines/indexing/db_schema.py
• wren-ai-service/src/pipelines/generation/question_recommendation.py
• wren-ai-service/src/pipelines/indexing/project_meta.py
• wren-ai-service/src/pipelines/generation/data_assistance.py
• wren-ai-service/src/pipelines/generation/chart_generation.py
• wren-ai-service/src/pipelines/generation/sql_tables_extraction.py

🧰 Additional context used

🧠 Learnings (10)

📓 Common learnings

Learnt from: wwwy3y3
PR: Canner/WrenAI#1585
File: wren-ui/src/pages/api/v1/generate_sql.ts:98-106
Timestamp: 2025-04-24T16:10:43.308Z
Learning: In the generate_sql API, allow users to specify language codes not predefined in the WrenAILanguage enum, passing them through directly rather than strictly validating against the enum.

Learnt from: cyyeh
PR: Canner/WrenAI#1112
File: wren-ai-service/src/pipelines/generation/sql_question.py:76-77
Timestamp: 2025-01-14T07:41:14.564Z
Learning: In the WrenAI codebase, extensive error handling for JSON parsing and array access operations in pipeline post-processing functions is not required, as indicated by the maintainer's preference to skip such suggestions.

wren-ai-service/src/pipelines/generation/intent_classification.py (1)

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

wren-ai-service/src/web/v1/services/ask.py (1)

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

wren-ai-service/src/pipelines/common.py (1)

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

wren-ai-service/src/pipelines/generation/sql_regeneration.py (2)

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

Learnt from: wwwy3y3
PR: Canner/WrenAI#1585
File: wren-ui/src/pages/api/v1/generate_sql.ts:98-106
Timestamp: 2025-04-24T16:10:43.308Z
Learning: In the generate_sql API, allow users to specify language codes not predefined in the WrenAILanguage enum, passing them through directly rather than strictly validating against the enum.

wren-ai-service/src/pipelines/generation/followup_sql_generation.py (2)

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

Learnt from: cyyeh
PR: Canner/WrenAI#1112
File: wren-ai-service/src/pipelines/generation/sql_question.py:76-77
Timestamp: 2025-01-14T07:41:14.564Z
Learning: In the WrenAI codebase, extensive error handling for JSON parsing and array access operations in pipeline post-processing functions is not required, as indicated by the maintainer's preference to skip such suggestions.

wren-ai-service/src/web/v1/services/question_recommendation.py (1)

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

wren-ai-service/src/pipelines/generation/utils/sql.py (2)

Learnt from: cyyeh
PR: Canner/WrenAI#1112
File: wren-ai-service/src/pipelines/generation/sql_question.py:76-77
Timestamp: 2025-01-14T07:41:14.564Z
Learning: In the WrenAI codebase, extensive error handling for JSON parsing and array access operations in pipeline post-processing functions is not required, as indicated by the maintainer's preference to skip such suggestions.

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (2)

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

Learnt from: cyyeh
PR: Canner/WrenAI#1112
File: wren-ai-service/src/pipelines/generation/sql_question.py:76-77
Timestamp: 2025-01-14T07:41:14.564Z
Learning: In the WrenAI codebase, extensive error handling for JSON parsing and array access operations in pipeline post-processing functions is not required, as indicated by the maintainer's preference to skip such suggestions.

wren-ai-service/src/pipelines/generation/sql_generation.py (1)

Learnt from: cyyeh
PR: Canner/WrenAI#1763
File: wren-ai-service/src/pipelines/generation/semantics_description.py:31-33
Timestamp: 2025-06-20T02:37:21.292Z
Learning: In the wren-ai-service codebase, when adding new fields like "alias" to the output of functions that use Pydantic models for validation, the user prefers not to update the corresponding Pydantic model definitions to include these new fields.

🧬 Code Graph Analysis (2)

wren-ai-service/src/pipelines/generation/intent_classification.py (1)

wren-ai-service/src/pipelines/common.py (1)

• build_table_ddl (29-63)

wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (2)

wren-ai-service/src/pipelines/generation/intent_classification.py (1)

• construct_db_schemas (235-265)

wren-ai-service/src/pipelines/common.py (1)

• build_table_ddl (29-63)

🔇 Additional comments (12)

wren-ai-service/src/pipelines/generation/utils/sql.py (2)

413-448: LGTM! Well-structured JSON field instructions.

The extraction of JSON-specific SQL instructions into a dedicated variable is a good refactoring that improves modularity. The instructions are comprehensive and cover the key aspects:

• Proper use of JSON_QUERY vs JSON_QUERY_ARRAY based on json_type
• Prohibition of deprecated functions
• Correct usage of LAX casting functions
• Proper UNNEST syntax for JSON arrays

This modular approach will enable conditional inclusion of these instructions based on the has_json_field flag mentioned in the AI summary.

---

250-250: LGTM! Simplified and clearer instruction.

The note about ORDER BY and LIMIT usage with UNION queries is now more concise and clear.

wren-ai-service/src/pipelines/generation/intent_classification.py (1)

262-262: LGTM! Correctly updated to handle new build_table_ddl signature.

The unpacking is properly updated to accommodate the new return signature that includes the has_json_field flag. The unused values are appropriately discarded with _ placeholders.

wren-ai-service/src/pipelines/generation/misleading_assistance.py (1)

30-30: LGTM! Standardized response length constraint.

Fixing the maximum response length to 100 words for all languages simplifies the logic and ensures consistent behavior. This is a good improvement over having language-specific constraints.

wren-ai-service/src/web/v1/services/ask.py (2)

485-485: LGTM! Properly integrated JSON field support in the ask method.

The has_json_field flag is correctly extracted from retrieval results and passed to the appropriate SQL generation pipelines (followup_sql_generation and sql_generation). This follows the same pattern as existing flags for calculated fields and metrics.

Also applies to: 500-500, 517-517

---

760-760: LGTM! Consistent JSON field support in ask_feedback method.

The has_json_field flag is properly integrated into the ask_feedback method, maintaining consistency with the ask method implementation. The flag is correctly passed to the sql_regeneration pipeline.

Also applies to: 785-785

wren-ai-service/src/pipelines/common.py (1)

31-31: LGTM! Well-implemented JSON field detection.

The addition of JSON field detection to build_table_ddl is implemented correctly:

1. Return signature updated: Properly changed from Tuple[str, bool] to Tuple[str, bool, bool] to include the new flag
2. Detection logic: Uses case-insensitive comparison column["data_type"].lower() == "json" which is consistent with existing patterns
3. Initialization: Properly initialized has_json_field = False
4. Return statement: Correctly includes the new flag in the tuple return

This change enables the JSON field awareness throughout the pipeline system and follows the same pattern as the existing has_calculated_field detection.

Also applies to: 34-34, 45-46, 56-63

wren-ai-service/src/pipelines/generation/sql_regeneration.py (1)

19-19: LGTM! Consistent JSON field awareness implementation.

The addition of json_field_instructions support follows the expected pattern and integrates well with the existing conditional instruction system. The removal of the configuration parameter appropriately simplifies the pipeline interface.

Also applies to: 59-61, 106-106, 120-120, 192-192, 207-207

wren-ai-service/src/web/v1/services/question_recommendation.py (1)

12-12: LGTM! Proper integration of JSON field awareness.

The service correctly extracts the has_json_field flag from retrieval results and propagates it to the SQL generation pipeline. The removal of Configuration usage aligns with the broader codebase cleanup.

Also applies to: 70-81, 105-105, 122-122

wren-ai-service/src/pipelines/generation/followup_sql_generation.py (1)

20-20: LGTM! Consistent implementation across SQL generation pipelines.

The JSON field awareness implementation follows the same pattern as other SQL generation pipelines, ensuring consistency across the codebase.

Also applies to: 49-51, 98-98, 112-112, 197-197, 221-221

wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (1)

241-241: LGTM! Comprehensive JSON field tracking implementation.

The implementation correctly tracks has_json_field throughout the schema retrieval process by leveraging the updated build_table_ddl function. The change from |= to explicit if statements improves code clarity while maintaining the same logic.

Also applies to: 245-245, 252-255, 286-286, 294-294, 358-358, 362-372, 405-405, 416-416

wren-ai-service/src/pipelines/generation/sql_generation.py (1)

19-19: LGTM! Completes the consistent JSON field awareness across SQL pipelines.

The implementation matches the pattern established in other SQL generation pipelines, ensuring uniform handling of JSON field instructions throughout the system.

Also applies to: 43-45, 94-94, 108-108, 193-193, 217-217

✨ Finishing Touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch chore/ai-service/improve-text2sql

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai auto-generate unit tests to generate unit tests for this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.