chore(wren-ai-service): improve ai service#1009
Conversation
cyyeh
added
module/ai-service
cyyeh
force-pushed
the
chore/ai-service/improve-text2sql
branch
9 times, most recently
from
ffcb484 to
b5f55b5
Compare
cyyeh
force-pushed
the
chore/ai-service/improve-text2sql
branch
from
0cc7ed0 to
90f1949
Compare
cyyeh
changed the title
|
@coderabbitai full review |
✅ Actions performedFull review triggered. |
WalkthroughThe pull request introduces several modifications across the Wren AI service, focusing on restructuring SQL processing modules, enhancing intent classification, and updating configuration handling. Key changes include the addition of a new target in the Changes
Sequence DiagramsequenceDiagram
participant User
participant IntentClassification
participant SQLGeneration
participant WrenEngine
User->>IntentClassification: Submit query
IntentClassification->>IntentClassification: Rephrase query
IntentClassification-->>User: Return rephrased query and reasoning
IntentClassification->>SQLGeneration: Classify intent
SQLGeneration->>WrenEngine: Generate and validate SQL
WrenEngine-->>SQLGeneration: Return SQL results
SQLGeneration-->>User: Provide query results
Poem
Tip CodeRabbit's docstrings feature is now available as part of our Early Access Program! Simply use the command 🪧 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: 1
🧹 Nitpick comments (13)
wren-ai-service/src/force_update_config.py (1)
4-21: Consider adding error handling for missing or malformed config documents.The logic assumes that
config.yamlhas valid YAML documents and that thepipelinedocument is always present. IfdocisNoneor fails to contain the required keys, the loop may silently skip them, creating a debugging challenge.Potential enhancements:
- Raise a custom error if pipeline documents are missing or do not contain the
"pipes"key.- Add validation for the YAML structure before writing back.
wren-ai-service/Justfile (1)
29-29: Recommended: Keep the 'start' dependency consistent with naming.Having
start: use-wren-ui-as-engineensures the config update runs automatically. However, if future changes require a different engine or approach, consider a naming convention that indicates exactly what is being done (e.g.,force-update-config).wren-ai-service/src/pipelines/generation/utils/sql.py (4)
16-16: Use module-level__name__for logger
Usinglogging.getLogger(__name__)instead of a hard-coded string can help identify the module automatically for debugging or log filtering across the codebase.
24-71: Solid error handling inrunmethod
- The method checks for missing
stepsand handles them gracefully with an early return.- The fallback to
descriptionin case of invalid or empty steps is a nice usage for user feedback.- One suggestion is to include more verbose logging when returning empty steps, which can aid in troubleshooting.
102-153: Robust parsing inSQLGenPostProcessor.run
- The conditional check for different reply data types covers a variety of payload structures, improving resilience.
- The nested try-except ensures partial results are still processed if there's an error in one of the elements.
- Consider adding unit tests for the mixed scenario of partially valid, partially invalid replies to verify resilience.
203-336: ExtensiveTEXT_TO_SQL_RULESguidelines
- This large monolithic string provides comprehensive SQL generation constraints, which is valuable, but can be hard to maintain.
- Consider splitting these rules into smaller segments or a structured documentation format for scalability.
- Ensure tests cover various rules, especially with special cases like JSON fields or time-based queries.
wren-ai-service/src/pipelines/generation/intent_classification.py (6)
43-43: Fix minor typo "rephrasedd"
Replace "rephrasedd" with "rephrased" for clarity.- - If the rephrasedd user's question is related to the previous question... + - If the rephrased user's question is related to the previous question...
45-49: Correct repeated "rephrasedd"
Lines 45 and 46 contain the same minor spelling error.- - The rephrasedd user's question involves specific data retrieval... - - The rephrasedd user's question references tables, columns, or specific data points... + - The rephrased user's question involves specific data retrieval... + - The rephrased user's question references tables, columns, or specific data points...
56-58: Resolve repeated "rephrasedd"
Please remove the extra "d."- - If the rephrasedd user's question is irrelevant to the given database schema... - - If the rephrasedd user's question is not related to the previous question... - - If the rephrasedd user's question contains SQL code. + - If the rephrased user's question is irrelevant to the given database schema... + - If the rephrased user's question is not related to the previous question... + - If the rephrased user's question contains SQL code.
60-64: Clean up "rephrasedd" occurrences
Lines 60, 61, 62, and 64 reference "the rephrasedd user's question."- - The rephrasedd user's question does not pertain to any aspect... - - The rephrasedd user's question might be a casual conversation... - - The rephrasedd user's question is vague... ... - - MUST explicitly add phrases from the rephrasedd user's question... + - The rephrased user's question does not pertain to any aspect... + - The rephrased user's question might be a casual conversation... + - The rephrased user's question is vague... ... + - MUST explicitly add phrases from the rephrased user's question...
71-72: Fix extra "d" in "rephrasedd"- - If the rephrasedd user's question is related to the previous question... - - If the rephrasedd user's question is related to the previous question... + - If the rephrased user's question is related to the previous question... + - If the rephrased user's question is related to the previous question...
76-77: Remove the extra "d" in line 77- - MUST explicitly add phrases from the rephrasedd user's question... + - MUST explicitly add phrases from the rephrased user's question...wren-ai-service/src/web/v1/services/ask.py (1)
199-205: Repeated pattern in constructing AskResultResponse
There are multiple similar blocks for settingAskResultResponse. Consider a helper function to reduce duplication.# Pseudo-refactor approach: # def build_default_response(status, the_type, rq, ir): # return AskResultResponse( # status=status, # type=the_type, # rephrased_question=rq, # intent_reasoning=ir, # )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
wren-ai-service/Justfile(2 hunks)wren-ai-service/demo/utils.py(0 hunks)wren-ai-service/src/__main__.py(1 hunks)wren-ai-service/src/force_update_config.py(1 hunks)wren-ai-service/src/pipelines/common.py(1 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(1 hunks)wren-ai-service/src/pipelines/generation/intent_classification.py(3 hunks)wren-ai-service/src/pipelines/generation/sql_breakdown.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_correction.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_expansion.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_regeneration.py(1 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/web/v1/services/ask.py(11 hunks)
💤 Files with no reviewable changes (1)
- wren-ai-service/demo/utils.py
🔇 Additional comments (38)
wren-ai-service/src/force_update_config.py (2)
23-28: Good practice to communicate success or failure.
Catching exceptions and printing a message is helpful for debugging. As a minor improvement, consider using structured logging or returning meaningful exit codes for CI/CD pipelines.
1-2: Check if yaml is in requirements.
Make sure that the yaml package (PyYAML) is pinned in your dependencies to prevent version mismatches that might cause problems with safe_load_all or safe_dump_all.
Please confirm that PyYAML is listed in your pyproject.toml or requirements.txt.
✅ Verification successful
Based on the search results, I can see that PyYAML is properly pinned in the project's dependencies. Specifically, in wren-ai-service/pyproject.toml:
pyyaml = "^6.0.2"
And this is also reflected in the detailed dependency lock file (poetry.lock) which contains the exact version and hashes for PyYAML 6.0.2.
PyYAML is properly pinned in your dependencies
The package is correctly pinned to version 6.0.2 in your project's dependencies, which ensures version consistency and prevents potential issues with safe_load_all or safe_dump_all functions.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Search for PyYAML in the project files.
rg "PyYAML|yaml"
Length of output: 30773
wren-ai-service/Justfile (1)
67-68: Well-designed target hook.
This new target seamlessly ties in the config update via poetry run python -m src.force_update_config. Ensure the config file is present before running, or handle the missing file scenario gracefully inside force_update_config.py.
wren-ai-service/src/__main__.py (1)
95-95: Great use of the reload_excludes parameter.
Excluding these folders from reload should speed up development. Confirm that no essential test or demo resources are inadvertently excluded from reload within the main runtime environment.
wren-ai-service/src/pipelines/generation/utils/sql.py (4)
1-4: Impressively clean setup of async and typed imports
The import statements (asyncio, logging, typing) and the straightforward organization establish a clear foundation. No issues stand out here.
19-23: Class declaration clarity
Defining a separate post-processor class (SQLBreakdownGenPostProcessor) aligns well with the single-responsibility principle and increases maintainability. Good job.
73-81: Readable CTE string construction
The _build_cte_query method composes the CTE clauses in a concise manner. Ensure thorough testing for special characters or edge cases in cte_name or sql.
82-100: Asynchronous SQL executability check
- The use of
aiohttp.ClientSession()and_engine.execute_sqlis straightforward. - Logging an exception with
logger.exceptionensures stack traces are captured. - Consider clarifying handling of partial successes or warnings if future expansions anticipate them.
wren-ai-service/src/pipelines/generation/sql_expansion.py (1)
14-15: Refactored imports for new SQL post-processor
The switch from legacy imports in src.pipelines.common to src.pipelines.generation.utils.sql import SQLGenPostProcessor centralizes SQL-related logic into one module, improving maintainability. Ensure you remove any unused imports from the old location to keep the codebase clean.
wren-ai-service/src/pipelines/generation/sql_generation.py (1)
14-15: Refined import paths for SQL generation
- Bringing in
show_current_timefromcommonis consistent with other modules, helping unify time-based context. - Importing the new
SQLGenPostProcessorfromutils.sqlfosters a more modular structure where all SQL-related components reside in one place.
wren-ai-service/src/pipelines/generation/followup_sql_generation.py (1)
14-15: Consistent usage of show_current_time and SQLGenPostProcessor
The updated imports align with the architecture changes. It’s good to see uniform patterns in follow-up generation, mirroring the main SQL generation logic.
wren-ai-service/src/pipelines/common.py (1)
2-2: Minimal type additions
The added Optional import suggests a small but helpful improvement in type hinting. Ensure all references in the file leverage this newly imported type for clarity.
wren-ai-service/src/pipelines/generation/sql_correction.py (1)
Line range hint 16-20: LGTM! Import path restructuring looks good.
The relocation of SQL-related components to utils.sql module improves code organization by grouping related functionality together.
Let's verify the new import path exists and contains all the required components:
✅ Verification successful
Import path exists but components need verification
Based on the code analysis, while the import path src.pipelines.generation.utils.sql exists and is used in multiple files, I need to verify if all the required components (TEXT_TO_SQL_RULES, SQLGenPostProcessor, sql_generation_system_prompt) are actually defined in the found sql.py file.
From the file content, I can confirm:
TEXT_TO_SQL_RULESis defined as a string constantSQLGenPostProcessoris defined as a component classsql_generation_system_promptis defined as a string constant
All three imported components are present in the source file with their expected definitions. The import path restructuring is valid and all components are properly defined in their new location.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the existence and contents of the new SQL utils module
# Test: Check if the new module exists and contains the required components
ast-grep --pattern 'from src.pipelines.generation.utils.sql import $_'
Length of output: 340
Script:
#!/bin/bash
# Let's verify if the required components exist in the utils.sql module
fd "sql.py" src/pipelines/generation/utils/ --exec cat {}
Length of output: 178
Script:
#!/bin/bash
# Let's try to locate the sql.py file first, then examine its contents
fd "sql.py" --exec cat {}
Length of output: 41525
wren-ai-service/src/pipelines/generation/sql_regeneration.py (1)
15-15: LGTM! Consistent with the module restructuring.
The import path change aligns with the codebase reorganization, maintaining consistency across the service.
wren-ai-service/src/pipelines/generation/sql_breakdown.py (1)
14-17: LGTM! Clean import restructuring with improved readability.
The multi-line import style using parentheses enhances code readability while maintaining consistency with the module reorganization.
wren-ai-service/src/pipelines/generation/intent_classification.py (6)
25-37: Enhance compliance with user data handling
These prompt instructions look generally sound. However, re-check compliance policies to ensure that any rephrasing of user queries does not inadvertently log or expose personally identifiable information (PII).
88-89: JSON output structure
These fields neatly align with the new rephrased question and reasoning logic.
250-250: Check for potential index error
classify_intent.get("replies")[0] may fail if “replies” is missing or empty. Although it’s wrapped in a try/except, consider verifying replies is non-empty beforehand.
252-254: Parsing keys from LLM output
Extracting "results", "rephrased_question", and "reasoning" is valid, and the exception pathway covers unexpected formats.
258-263: Fallback classification logic
Returning "TEXT_TO_SQL" when an exception occurs might incorrectly convey the final intent to upstream consumers. Confirm this meets your user experience requirements.
271-272: Added fields to IntentClassificationResult
Including rephrased_question and reasoning is consistent with the prompt updates and final usage.
wren-ai-service/src/web/v1/services/ask.py (17)
93-94: Added optional fields to AskResultResponse
This supports the expanded classification details without breaking existing responses.
142-145: Initialize rephrased_question and intent_reasoning
Defines these fields at the start, ensuring they remain in scope even if classification fails.
162-166: Extract classification result fields
Retrieving rephrased_question and reasoning directly from the pipeline’s output is straightforward and clearly handled.
167-170: Use rephrased question as fallback
Switching to the rephrased question ensures user experience is streamlined, though confirm it doesn’t disrupt any logs or analytics.
175-176: Include context in MISLEADING_QUERY response
Populating rephrased_question and intent_reasoning can aid in debugging and user messaging.
183-183: Pass user_query to data_assistance
Adopting the possibly rephrased query fosters consistency across all pipelines.
192-195: Set final response when intent is GENERAL
Enhances clarity by returning rephrased_question and intent_reasoning in the final result.
210-212: Same response-building pattern
216-216: Use user_query in retrieval
Ensures consistent input to the retrieval pipeline, avoiding mismatch with rephrased data.
223-223: Mindful logging of user_query
At line 223, indefinite logging of user_query may risk exposing personal data. Lines 232-233 consistently set rephrased details in the error response, which is helpful.
Do you want to partially redact or anonymize the query if it contains identifiable info?
Also applies to: 232-233
242-244: Repeated response construction
248-248: Historical question pipeline
Consistently using user_query fosters alignment with the rephrased approach.
274-274: Follow-up SQL generation
Carrying forward user_query is consistent and streamlined.
284-284: SQL generation pipeline
Again, using user_query maintains continuity among pipeline steps.
338-339: Final response includes rephrased info
Exposing rephrased_question and intent_reasoning in the results can assist with user comprehension.
344-344: Potential PII in logs
Line 344 logs user_query, which might present privacy issues. Lines 353-354 preserve rephrased/meta data for debugging. Check for potential PII exposures.
Also applies to: 353-354
363-363: Include contextual fields in error handling
Carrying rephrased_question and intent_reasoning into error scenarios improves visibility into failed queries.
Also applies to: 370-371
| async def _classify_invalid_generation_results( | ||
| self, generation_results: List[Dict[str, str]], project_id: str | None = None | ||
| ) -> List[Optional[Dict[str, str]]]: | ||
| valid_generation_results = [] | ||
| invalid_generation_results = [] | ||
|
|
||
| async def _task(result: Dict[str, str]): | ||
| quoted_sql, no_error = add_quotes(result["sql"]) | ||
|
|
||
| if no_error: | ||
| status, _, addition = await self._engine.execute_sql( | ||
| quoted_sql, session, project_id=project_id | ||
| ) | ||
|
|
||
| if status: | ||
| valid_generation_results.append( | ||
| { | ||
| "sql": quoted_sql, | ||
| "correlation_id": addition.get("correlation_id", ""), | ||
| } | ||
| ) | ||
| else: | ||
| invalid_generation_results.append( | ||
| { | ||
| "sql": quoted_sql, | ||
| "type": "DRY_RUN", | ||
| "error": addition.get("error_message", ""), | ||
| "correlation_id": addition.get("correlation_id", ""), | ||
| } | ||
| ) | ||
| else: | ||
| invalid_generation_results.append( | ||
| { | ||
| "sql": result["sql"], | ||
| "type": "ADD_QUOTES", | ||
| "error": "add_quotes failed", | ||
| } | ||
| ) | ||
|
|
||
| async with aiohttp.ClientSession() as session: | ||
| tasks = [ | ||
| _task(generation_result) for generation_result in generation_results | ||
| ] | ||
| await asyncio.gather(*tasks) | ||
|
|
||
| return valid_generation_results, invalid_generation_results | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Concurrent classification of SQL generation results
- The
_taskinner function and usage ofasyncio.gatheris a clear concurrency pattern. - The classification logic for
invalid_generation_resultsis robust, though including more detailed error codes could improve debugging. - Currently, there is no explicit limit on concurrency or upon large input. For extensive inputs, controlling concurrency with a semaphore might be beneficial.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
docker/config.example.yaml (1)
11-11: LGTM! Consider documenting the seed parameter's purpose.The addition of
seed: 0to both models is a good practice for ensuring deterministic outputs, which aligns well with the existing deterministic configuration (temperature: 0, n: 1).Consider adding a comment in the configuration files or updating the documentation to explain:
- The purpose of the seed parameter
- Why the value 0 was chosen
- How this affects the model's output determinism
Also applies to: 21-21
wren-ai-service/tools/config/config.full.yaml (1)
11-11: LGTM! Complete configuration coverage achieved.The seed parameter is consistently applied across all configuration variants, including the full configuration with additional embedder options.
Since this change introduces deterministic behavior:
- Consider adding this configuration detail to the architecture documentation
- Document how this change affects the reproducibility of results across different deployment environments
Also applies to: 21-21
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
deployment/kustomizations/base/cm.yaml(1 hunks)docker/config.example.yaml(2 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(2 hunks)wren-ai-service/tools/config/config.example.yaml(2 hunks)wren-ai-service/tools/config/config.full.yaml(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/src/pipelines/generation/sql_generation.py
🔇 Additional comments (4)
docker/config.example.yaml (1)
Line range hint 4-24: Verify deterministic behavior impact on pipeline components.
The seed parameter will affect all pipeline components using these models. Please verify that deterministic outputs are desired for all the following components:
- db_schema_retrieval
- sql_generation
- sql_correction
- followup_sql_generation
- sql_summary
- sql_answer
- sql_breakdown
- sql_expansion
- sql_explanation
- sql_regeneration
- semantics_description
- relationship_recommendation
- question_recommendation
- intent_classification
- data_assistance
- preprocess_sql_data
- chart_generation
- chart_adjustment
✅ Verification successful
Deterministic behavior is consistently configured across all pipeline components
The search results confirm that:
- All pipeline components are using the same model configurations with deterministic settings (seed=0)
- The configuration is consistently applied across all environments:
- Development (docker/config.example.yaml)
- Testing (wren-ai-service/tools/config/config.*.yaml)
- Production (deployment/kustomizations/base/cm.yaml)
- The deterministic behavior is intentionally set for all components that use these models, including:
- Schema understanding (db_schema_retrieval)
- SQL operations (generation, correction, explanation, etc.)
- Natural language processing (intent_classification, data_assistance)
- Visualization (chart_generation, chart_adjustment)
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Find all pipeline components using the affected models
# Expected: All components should be reviewed for deterministic behavior requirements
# Search for pipeline components using the affected models
rg "litellm_llm\.gpt-4o-(?:mini-2024-07-18|2024-08-06)" -A 3 -B 3
# Search for any documentation or comments about deterministic behavior requirements
rg -i "deterministic|reproducible|seed" --type yaml --type md
Length of output: 18096
wren-ai-service/tools/config/config.example.yaml (1)
11-11: LGTM! Configuration consistency maintained.
The seed parameter is consistently applied across docker and development environments, ensuring the same deterministic behavior in both contexts.
Also applies to: 21-21
deployment/kustomizations/base/cm.yaml (2)
62-62: Ensure consistent model behavior with fixed seed.
The addition of seed: 0 to both model configurations will ensure deterministic outputs, which is beneficial for testing and reproducibility. This is a good practice for maintaining consistent behavior across model invocations.
Let's verify this configuration is consistent across other files:
Also applies to: 72-72
✅ Verification successful
Model configuration is consistently set across all files
The verification confirms that the seed: 0 parameter is consistently set for both models (gpt-4o-mini-2024-07-18 and gpt-4o-2024-08-06) across all configuration files in the repository:
docker/config.example.yamlwren-ai-service/tools/config/config.full.yamlwren-ai-service/tools/config/config.example.yamldeployment/kustomizations/base/cm.yaml
All configurations maintain identical settings with temperature: 0, n: 1, seed: 0, and max_tokens: 4096, ensuring consistent deterministic behavior across different environments.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if seed parameter is consistently set across all configuration files
# Search for model configurations in yaml files
rg --type yaml -A 10 "model: gpt-4o-(mini-2024-07-18|2024-08-06)" | grep -A 5 "kwargs:"
Length of output: 2824
Line range hint 89-166: Review pipeline configurations for model assignments.
The pipeline configuration appears to consistently use gpt-4o-mini-2024-07-18 across various tasks. Consider if any high-complexity tasks (like sql_generation or intent_classification) would benefit from using the potentially more capable gpt-4o-2024-08-06 model instead.
Let's analyze the pipeline usage patterns:
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wren-ai-service/tests/pytest/test_usecases.py (2)
43-43: Guard against potential KeyError
question["question"]assumes that each item inquestionshas a"question"key. Consider providing defaults or validations to avoidKeyErrorif the key is missing.
47-47: Nitpick: Revise exception message for clarityThe exception message might appear ambiguous to users encountering a missing file. Rewording for clarity could better direct them toward missing or incorrect YAML or JSON file paths.
- f"tests/data/usecases/{usecase}/mdl.json or tests/data/usecases/{usecase}/questions.yaml not found" + f"Could not find mdl.json or questions.yaml for the '{usecase}' usecase. Please verify file names and paths."
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/tests/pytest/test_usecases.py(2 hunks)
🔇 Additional comments (2)
wren-ai-service/tests/pytest/test_usecases.py (2)
13-13: Use yaml.safe_load responsibly
Good practice using yaml.safe_load; ensure that PyYAML is included in your project's dependencies to avoid import errors in production or when running tests.
38-39: Validate YAML file structure
If the questions YAML file is improperly formatted (e.g., missing question keys), the test will fail at runtime. Consider adding validation logic or error handling for malformed or incomplete YAML files.
2 participants
Summary by CodeRabbit
New Features
use-wren-ui-as-engineto streamline application startup.rephrased_questionandintent_reasoninginAskResultResponsefor enhanced query context.update_configto update YAML configuration for pipeline entries.SQLBreakdownGenPostProcessorandSQLGenPostProcessor, for improved SQL generation processing.Bug Fixes
askmethod to preserve additional context.Refactor
Chores
seedparameter for randomness control.