Repository navigation
chore(wren-ai-service): improve text2sql - #1208
Conversation
cyyeh
added
module/ai-service
WalkthroughThe pull request introduces significant modifications to the Wren AI service's pipeline architecture, focusing on enhancing flexibility and data handling across multiple components. The changes primarily involve updating function signatures, modifying how retrieval results are processed, and introducing new parameters to track calculated fields and metrics. These modifications span several files in the AI service, affecting SQL generation, retrieval, and evaluation pipelines. Changes
Sequence DiagramsequenceDiagram
participant User
participant RetrievalPipeline
participant SQLGeneration
participant AskPipeline
User->>RetrievalPipeline: Submit Query
RetrievalPipeline-->>RetrievalPipeline: Check DB Schemas
RetrievalPipeline->>SQLGeneration: Pass Retrieval Results
SQLGeneration-->>SQLGeneration: Generate SQL
SQLGeneration->>AskPipeline: Return Results with Flags
AskPipeline-->>AskPipeline: Process Results
AskPipeline->>User: Return Final Answer
Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
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)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
wren-ai-service/src/pipelines/generation/intent_classification.py (1)
210-211: Consider utilizing the calculated field flag.The code discards the boolean flag indicating the presence of calculated fields returned by
build_table_ddl. This information could be valuable for intent classification.Consider updating the code to utilize this information:
- ddl, _ = build_table_ddl(table_schema) - db_schemas_in_ddl.append(ddl) + ddl, has_calculated_field = build_table_ddl(table_schema) + db_schemas_in_ddl.append({ + "ddl": ddl, + "has_calculated_field": has_calculated_field + })wren-ai-service/src/pipelines/generation/utils/sql.py (1)
668-686: Enhance fiscal year instruction formatting.The fiscal year instruction could be more explicit and better formatted for clarity.
Consider this improvement:
- instructions += f"\n- For calendar year related computation, it should be started from {configuration.fiscal_year.start} to {configuration.fiscal_year.end}\n\n" + instructions += f""" +#### Fiscal Year Configuration #### +For calendar year computations: +- Start date: {configuration.fiscal_year.start} +- End date: {configuration.fiscal_year.end} + +"""
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
wren-ai-service/eval/pipelines.py(7 hunks)wren-ai-service/eval/prediction.py(1 hunks)wren-ai-service/src/core/engine.py(0 hunks)wren-ai-service/src/pipelines/common.py(3 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(8 hunks)wren-ai-service/src/pipelines/generation/intent_classification.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_correction.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(6 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(8 hunks)wren-ai-service/src/pipelines/retrieval/retrieval.py(5 hunks)wren-ai-service/src/web/v1/services/ask.py(4 hunks)wren-ai-service/src/web/v1/services/question_recommendation.py(1 hunks)wren-ai-service/tests/pytest/pipelines/generation/test_ask.py(0 hunks)
💤 Files with no reviewable changes (2)
- wren-ai-service/src/core/engine.py
- wren-ai-service/tests/pytest/pipelines/generation/test_ask.py
🔇 Additional comments (29)
wren-ai-service/eval/prediction.py (1)
88-88: LGTM! Help text is now more accurate.The updated help text better reflects that the argument accepts a full file path rather than just a filename in a specific folder.
wren-ai-service/src/pipelines/generation/sql_correction.py (1)
87-87: LGTM! Improved naming convention and simplified response structure.The renaming from
CorrectedSQLResulttoSqlCorrectionResultaligns with the codebase's naming conventions, and the response structure has been simplified by removing the nested results wrapper.Let's verify that the removal of the
CorrectedResultsclass doesn't break any existing code:Also applies to: 92-92
✅ Verification successful
✅ Verified: Class renaming and structure simplification are safe
The codebase shows no remaining references to the old
CorrectedResultsorCorrectedSQLResultclasses. The newSqlCorrectionResultclass is properly defined and used only within its module, confirming that the renaming and structure simplification are safe changes.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Search for any remaining references to CorrectedResults rg "CorrectedResults" -lLength of output: 24
Script:
#!/bin/bash # Search for various case combinations and both class names echo "=== Searching for old class name variations ===" rg -i "correctedsqlresult|correctedresults" --type py echo -e "\n=== Searching for new class name ===" rg "SqlCorrectionResult" --type py -A 2 echo -e "\n=== Checking import statements ===" rg "^from.*import.*Corrected|^import.*Corrected" --type pyLength of output: 891
wren-ai-service/src/pipelines/generation/sql_generation.py (2)
61-62: LGTM! Well-structured parameter propagation.The new parameters
has_calculated_fieldandhas_metricare consistently propagated through the pipeline, from the function signature to the execution inputs.Also applies to: 133-134, 145-146
30-33: LGTM! Enhanced prompt template organization.The addition of a conditional instructions section improves the prompt structure and clarity.
wren-ai-service/src/pipelines/generation/followup_sql_generation.py (2)
74-75: LGTM! Consistent with sql_generation.py changes.The addition of
has_calculated_fieldandhas_metricparameters maintains consistency with the changes in sql_generation.py.Also applies to: 150-151, 163-164
50-56: LGTM! Improved context organization.The context section now clearly separates and presents previous SQL summary and query information, enhancing readability and context understanding.
wren-ai-service/src/web/v1/services/question_recommendation.py (2)
76-79: LGTM! Improved retrieval result handling.The code now correctly extracts retrieval results and flags from the nested dictionary structure, with proper default values.
85-86: Verify SQL generation with new flags.The new flags are correctly passed to the SQL generation pipeline, enhancing its ability to handle calculated fields and metrics.
Run this script to verify the SQL generation pipeline handles these flags correctly:
✅ Verification successful
SQL generation properly handles calculated field and metric flags
The flags are correctly implemented throughout the entire pipeline:
- Detected during schema analysis
- Propagated through retrieval and service layers
- Used in both SQL generation and followup SQL generation
- Properly influence SQL generation instructions
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify SQL generation pipeline's handling of calculated fields and metrics flags # Test: Search for usage of has_calculated_field and has_metric parameters in SQL generation rg -A 5 'has_calculated_field|has_metric' 'src/pipelines/generation/'Length of output: 135
Script:
#!/bin/bash # First, find all Python files in the wren-ai-service directory fd -e py . wren-ai-service | while read -r file; do echo "=== Checking $file ===" rg -A 5 'has_calculated_field|has_metric' "$file" || true doneLength of output: 50517
wren-ai-service/eval/pipelines.py (6)
33-36: LGTM! Enhanced model deployment.The function now correctly handles multiple pipeline components using
asyncio.gather.
195-197: LGTM! Improved retrieval result extraction.The code now correctly extracts retrieval results from the nested dictionary structure.
248-249: LGTM! Added support for calculated fields and metrics.The code now correctly passes the flags to enhance SQL generation.
274-277: LGTM! Enhanced accuracy metric configuration.The code now correctly configures the accuracy metric with semantics comparison support.
299-306: LGTM! Improved indexing initialization.The code now correctly initializes both DB schema and table description indexing.
329-333: LGTM! Enhanced retrieval result processing.The code now correctly:
- Extracts retrieval results and flags from the nested structure
- Passes these flags to the SQL generation pipeline
- Updates the prediction dictionary with the flags
Also applies to: 338-339, 344-345
wren-ai-service/src/pipelines/retrieval/retrieval.py (9)
223-224: LGTM! Added flags for tracking special fields.The code now correctly initializes flags to track calculated fields and metrics.
228-230: LGTM! Enhanced table DDL handling.The code now correctly:
- Unpacks the tuple returned by
build_table_ddl- Updates the calculated field flag based on the result
237-237: LGTM! Added metric tracking.The code now correctly sets the metric flag when processing metric tables.
246-247: LGTM! Enhanced return value structure.The code now correctly includes calculated field and metric flags in the return value.
Also applies to: 253-254
271-273: LGTM! Updated DDL building.The code now correctly unpacks the tuple returned by
build_table_ddl.
305-305: LGTM! Updated return type annotation.The function's return type is now correctly annotated as
dict[str, Any].
319-320: LGTM! Enhanced table schema processing.The code now correctly:
- Tracks calculated fields and metrics
- Processes table schemas with specific columns
- Updates flags based on the results
Also applies to: 324-329, 331-332
340-340: LGTM! Added metric tracking in retrieval results.The code now correctly sets the metric flag when processing metric tables.
344-348: LGTM! Enhanced return value structure.The code now correctly returns a dictionary containing:
- Retrieval results
- Calculated field flag
- Metric flag
Also applies to: 352-358
wren-ai-service/src/web/v1/services/ask.py (2)
245-248: LGTM! Improved dictionary access pattern.The change enhances code robustness by safely accessing nested dictionary values and improves maintainability by storing intermediate results.
296-297: LGTM! Consistent parameter passing.The new parameters enhance SQL generation by providing additional context about calculated fields and metrics, and they are used consistently across both code paths.
Also applies to: 308-309
wren-ai-service/src/pipelines/generation/utils/sql.py (2)
8-8: LGTM! Improved result handling and type safety.The changes simplify the code by:
- Directly accessing SQL from replies without unnecessary data transformation
- Updating type hints to reflect the actual data structure
- Improving error handling by using raw SQL in error messages
Also applies to: 122-122, 127-127, 129-129, 152-152, 157-158, 184-184
333-341: LGTM! Enhanced SQL generation prompt.The updated prompt improves SQL generation quality by requiring explicit reasoning about queries before generating SQL. This structured approach helps ensure more accurate and contextual SQL generation.
wren-ai-service/src/pipelines/common.py (2)
1-1: LGTM! Improved type safety.The changes enhance type safety by properly annotating the tuple return type and organizing imports correctly.
Also applies to: 10-10
12-13: LGTM! Clean implementation of calculated field detection.The changes effectively track the presence of calculated fields while maintaining the original DDL generation functionality. The implementation is clean and non-intrusive.
Also applies to: 17-18, 33-33
Uh oh!
There was an error while loading. Please reload this page.