chore(wren-ai-service): improve text2sql (ai-env-changed) - #1936
Conversation
WalkthroughAdds a new SQL diagnosis generation pipeline and integrates it into Ask and AskFeedback flows; exposes SQLDiagnosis in exports; updates sql_correction to accept sql_functions; preserves original_sql across invalid results; adjusts engine error payloads; adds allow_sql_diagnosis config flag and pipeline entries in runtime/docs configs. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant AskSvc as AskService
participant Gen as SQLGenerator
participant Exec as SQLDryRun
participant Diag as SQLDiagnosis
participant Corr as SQLCorrection
User->>AskSvc: Ask question
AskSvc->>Gen: generate SQL
Gen-->>AskSvc: generation_result
AskSvc->>Exec: dry-run SQL
Exec-->>AskSvc: error_response (error_message, error_sql?)
alt error_response present
AskSvc->>Diag: contexts (schema), original_sql, invalid_sql, error_message, language
Diag-->>AskSvc: {reasoning, can_be_corrected}
alt can_be_corrected == true
AskSvc->>Corr: contexts, invalid_generation_result{sql: original_sql, error: reasoning}, sql_functions
Corr-->>AskSvc: {valid_sql | invalid_sql}
AskSvc->>Exec: dry-run corrected SQL (if valid)
Exec-->>AskSvc: {ok | error}
else cannot be corrected
AskSvc-->>User: report diagnostic reasoning / stop retries
end
else no error
AskSvc-->>User: return results
end
sequenceDiagram
autonumber
actor User
participant Feedback as AskFeedbackService
participant Diag as SQLDiagnosis
participant Corr as SQLCorrection
User->>Feedback: feedback with invalid SQL
alt invalid SQL
Feedback->>Diag: table_ddls, original_sql, invalid_sql, error_message, language
Diag-->>Feedback: {reasoning, can_be_corrected}
alt can_be_corrected == true
Feedback->>Corr: contexts, invalid_generation_result{sql: original_sql, error: reasoning}, sql_functions, project_id
Corr-->>Feedback: {valid | invalid}
else cannot be corrected
Feedback-->>User: return original invalid SQL/error
end
else valid
Feedback-->>User: proceed normally
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
cyyeh
added
module/ai-service
github-actions
Bot
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
wren-ai-service/src/pipelines/generation/sql_correction.py (2)
110-126: Guard against missing/empty replies to prevent runtime errors.If the generator returns no replies, this will raise at post-processing time. Add a defensive check.
async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: - return await post_processor.run( - generate_sql_correction.get("replies"), + replies = generate_sql_correction.get("replies") or [] + if not replies: + raise ValueError("SQLCorrection: generator returned no replies") + return await post_processor.run( + replies, project_id=project_id, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, )
74-79: Ensureinvalid_generation_resultalways includes.error
The prompt template insql_correction.py(lines 75–76) reads{{ invalid_generation_result.error }}, but inwren-ai-service/src/pipelines/generation/utils/sql.pyall theinvalid_generation_resultbranches only set keys like"sql","original_sql","type"or use"error_sql"—there’s no"error"key. Add an"error"field (e.g. fromerror_messageoraddition.get("error_message")) in every branch that buildsinvalid_generation_result.
🧹 Nitpick comments (19)
wren-ai-service/docs/config_examples/config.qwen3.yaml (1)
186-187: Consider a JSON-strong LLM alias for structured outputs.If
sql_diagnosisexpects JSON, prefer the alias configured withresponse_format: json_object(you already definedqwen3-32b).- - name: sql_diagnosis - llm: litellm_llm.default + - name: sql_diagnosis + llm: litellm_llm.qwen3-32bdocker/config.example.yaml (1)
164-165: Good addition; keep ordering consistent with examples.
sql_diagnosisfollowssql_tables_extractionas expected.If the diagnosis step emits structured fields, consider pointing it to an alias with deterministic JSON mode in this example to guide users.
wren-ai-service/tools/config/config.full.yaml (1)
175-176: LGTM; mirrors docker example.Placement and wiring look correct.
Optionally document in this file that
sql_diagnosisbenefits from an LLM/alias configured for JSON responses.wren-ai-service/docs/config_examples/config.groq.yaml (1)
147-148: Consistent addition; confirm model supports desired output format.Groq model wiring is fine. If
sql_diagnosisrelies on JSON output, addresponse_format: { type: json_object }to the chosen model or select an alias configured for JSON.wren-ai-service/docs/config_examples/config.azure.yaml (1)
156-157: Clarify coexistence with sql_tables_extraction.You’ve added sql_diagnosis while sql_tables_extraction remains defined above. If sql_diagnosis supersedes the old step in your runtime pipelines, consider removing sql_tables_extraction from example configs to avoid confusion. Otherwise, add a brief comment explaining when each is used.
wren-ai-service/src/providers/engine/wren.py (1)
89-102: Good capture of error_sql; unify correlation_id extraction and log it.Nice addition pulling dialectSql/plannedSql into error_sql. To improve observability and consistency, extract correlation_id the same way for success and error paths and include it in the error log.
Apply this diff within execute_sql’s error return and log:
- logger.error(f"Error executing SQL: {error_message}") + # include correlation id when available + logger.error( + f"Error executing SQL: {error_message} " + f"(correlation_id={res_json.get('errors', [{}])[0].get('extensions', {}).get('other', {}).get('correlationId', res_json.get('correlationId', ''))})" ) @@ - "correlation_id": ( - res_json.get("errors", [{}])[0] - .get("extensions", {}) - .get("other", {}) - .get("correlationId") - ), + "correlation_id": res_json.get("errors", [{}])[0] + .get("extensions", {}) + .get("other", {}) + .get("correlationId", res_json.get("correlationId", "")),Additionally, mirror this correlation_id fallback in the success-path returns (not shown in this hunk) so all paths use the same logic.
Outside this hunk, consider a small helper to DRY it up:
def _cid(j: dict) -> str: return ( j.get("errors", [{}])[0] .get("extensions", {}) .get("other", {}) .get("correlationId") or j.get("correlationId", "") or "" )Then use "correlation_id": _cid(res_json) everywhere.
Also applies to: 111-115
wren-ai-service/src/globals.py (2)
152-153: Verify intended order of sql_diagnosis in AskService.It currently runs after sql_functions_retrieval (and after sql_correction earlier in the dict). If the intention is to diagnose before correction, move it earlier; otherwise add a short comment indicating it is a post-correction diagnostic.
Example reordering if you intended pre-correction:
- "sql_correction": _sql_correction_pipeline, - "followup_sql_generation": generation.FollowUpSQLGeneration( - **pipe_components["followup_sql_generation"], - ), - "sql_functions_retrieval": _sql_functions_retrieval_pipeline, - "sql_diagnosis": _sql_diagnosis_pipeline, + "sql_diagnosis": _sql_diagnosis_pipeline, + "sql_correction": _sql_correction_pipeline, + "followup_sql_generation": generation.FollowUpSQLGeneration( + **pipe_components["followup_sql_generation"], + ), + "sql_functions_retrieval": _sql_functions_retrieval_pipeline,
172-173: AskFeedbackService: confirm diagnosis placement.Same concern: sql_diagnosis comes after sql_correction. If the goal is to surface actionable diagnostics before correction, swap their order.
- "sql_correction": _sql_correction_pipeline, - "sql_diagnosis": _sql_diagnosis_pipeline, + "sql_diagnosis": _sql_diagnosis_pipeline, + "sql_correction": _sql_correction_pipeline,wren-ai-service/docs/config_examples/config.zhipu.yaml (1)
194-195: Keep examples focused: remove legacy step or document both.You added sql_diagnosis but still list sql_tables_extraction. If only one is expected in typical setups, prune the other; otherwise annotate when to use each to prevent misconfiguration.
wren-ai-service/docs/config_examples/config.ollama.yaml (1)
146-147: Same as above—avoid dual steps unless both are required.Consider removing sql_tables_extraction or adding a brief note clarifying why both appear and how they’re used by services.
wren-ai-service/docs/config_examples/config.anthropic.yaml (1)
143-144: Confirm sql_diagnosis placement and remove legacy step if unused.If
sql_tables_extractionis deprecated in favor ofsql_diagnosis, consider removing it here to avoid confusion and unnecessary pipeline registrations. Otherwise, document why both are needed and their intended order.wren-ai-service/src/web/v1/services/ask.py (3)
518-533: Short-circuit is fine; add a minimal log for diagnoser verdict.A tiny debug log before
breakwill help triage non-syntax failures observed by users.- if not is_sql_syntax_issue: - break + if not is_sql_syntax_issue: + logger.debug("sql_diagnosis: non-syntax issue detected; skipping correction") + break
498-505: Don’t charge a retry when no correction is attempted.Increment retries only when proceeding to
sql_correctionto preserve attempts for genuine corrections.- current_sql_correction_retries += 1 + # increment only when a correction attempt will be made (see below)And after confirming
is_sql_syntax_issue:- if not is_sql_syntax_issue: - break + if not is_sql_syntax_issue: + break + current_sql_correction_retries += 1
596-603: Surface TIME_OUT distinctly for better UX.When prior loop breaks on TIME_OUT,
error_messagemay be None and users only see “No relevant SQL”. Consider propagating a clearer timeout message.- message=error_message or "No relevant SQL", + message=error_message or "SQL validation timed out or no relevant SQL",wren-ai-service/src/web/v1/services/ask_feedback.py (1)
201-208: Add a small debug crumb for diagnoser output.Helps correlate user-facing failures with diagnoser decisions.
- sql_diagnosis_results = await self._pipelines[ + sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( contexts=table_ddls, original_sql=original_sql, invalid_sql=invalid_sql, error_message=error_message, ) + logger.debug("sql_diagnosis completed for feedback path")wren-ai-service/src/pipelines/generation/sql_correction.py (2)
33-36: Fix prompt typos and tighten wording.Minor text issues in the instructions (e.g., “firgure”). Tighten phrasing for clarity.
-1. First, think hard about the error message, and firgure out the root cause first. -2. Then, generate the reasoning behind the correction. -3. Finally, generate the syntactically correct ANSI SQL query based on the reasoning to correct the error. -4. You could try to use other methods(new functions, etc.) to rewrite the SQL query to correct the error, but you should not change the original semantics of the SQL query. +1. Carefully analyze the error message and identify the root cause. +2. Provide concise reasoning for the correction. +3. Produce a syntactically correct ANSI SQL query based on that reasoning. +4. You may use alternative methods (e.g., new functions) to correct the SQL, but do not change its original semantics.
103-108: Return type annotation mismatches actual return (trace_cost expects a tuple).The function returns
(result, generator_name)before the decorator unwraps it. Annotate accordingly to avoid confusion with type checkers.-async def generate_sql_correction( - prompt: dict, generator: Any, generator_name: str -) -> dict: +async def generate_sql_correction( + prompt: dict, generator: Any, generator_name: str +) -> tuple[dict, str]:wren-ai-service/src/pipelines/generation/sql_diagnosis.py (2)
27-31: Polish diagnosis instructions for clarity.-1. First, think hard about the error message, and analyze the invalid SQL query to figure out the root cause and which part is incorrect. -2. Then, map the incorrect part of the invalid SQL query to the corresponding part of the original SQL query. -3. Then, return the reasoning behind the diagnosis.(You should give me the part of the original SQL query that is incorrect and the reason why it is incorrect) -4. Also, return a boolean value to indicate whether the issue is a SQL syntax issue. +1. Carefully analyze the error message and the invalid SQL to identify the root cause and the incorrect part. +2. Map the incorrect part of the invalid SQL to the corresponding part of the original SQL. +3. Provide concise reasoning for the diagnosis (include the specific part of the original SQL and why it is incorrect). +4. Indicate whether the issue is a SQL syntax issue with a boolean flag.
78-84: Return type annotation mismatches actual return (trace_cost expects a tuple).Align the type hint with the pre-decorator return.
-async def generate_sql_diagnosis( - prompt: dict, generator: Any, generator_name: str -) -> dict: +async def generate_sql_diagnosis( + prompt: dict, generator: Any, generator_name: str +) -> tuple[dict, str]:
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (25)
deployment/kustomizations/base/cm.yaml(1 hunks)docker/config.example.yaml(1 hunks)wren-ai-service/docs/config_examples/config.anthropic.yaml(1 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.bedrock.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_vertexai.yaml(1 hunks)wren-ai-service/docs/config_examples/config.grok.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.lm_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/docs/config_examples/config.open_router.yaml(1 hunks)wren-ai-service/docs/config_examples/config.qwen3.yaml(1 hunks)wren-ai-service/docs/config_examples/config.zhipu.yaml(1 hunks)wren-ai-service/src/globals.py(3 hunks)wren-ai-service/src/pipelines/generation/__init__.py(2 hunks)wren-ai-service/src/pipelines/generation/sql_correction.py(7 hunks)wren-ai-service/src/pipelines/generation/sql_diagnosis.py(1 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(3 hunks)wren-ai-service/src/providers/engine/wren.py(1 hunks)wren-ai-service/src/web/v1/services/ask.py(2 hunks)wren-ai-service/src/web/v1/services/ask_feedback.py(1 hunks)wren-ai-service/tools/config/config.example.yaml(1 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-03-18T10:28:10.593Z
Learnt from: andreashimin
PR: Canner/WrenAI#1414
File: wren-ui/src/apollo/server/utils/error.ts:0-0
Timestamp: 2025-03-18T10:28:10.593Z
Learning: The typo in error code enum names ("IDENTIED_AS_GENERAL" and "IDENTIED_AS_MISLEADING_QUERY" instead of "IDENTIFIED_AS_GENERAL" and "IDENTIFIED_AS_MISLEADING_QUERY") in wren-ui/src/apollo/server/utils/error.ts was recognized but intentionally not fixed as it was considered out of scope for PR #1414.
Applied to files:
wren-ai-service/src/providers/engine/wren.py
🧬 Code graph analysis (6)
wren-ai-service/src/pipelines/generation/__init__.py (1)
wren-ai-service/src/pipelines/generation/sql_diagnosis.py (1)
SQLDiagnosis(112-152)
wren-ai-service/src/pipelines/generation/sql_correction.py (3)
wren-ai-service/src/pipelines/generation/utils/sql.py (2)
construct_instructions(480-489)run(29-69)wren-ai-service/src/pipelines/retrieval/sql_functions.py (1)
SqlFunction(20-47)wren-ai-service/src/pipelines/generation/sql_generation.py (1)
run(177-219)
wren-ai-service/src/web/v1/services/ask.py (2)
wren-ai-service/src/pipelines/generation/sql_correction.py (1)
run(176-206)wren-ai-service/src/pipelines/generation/sql_generation.py (1)
run(177-219)
wren-ai-service/src/web/v1/services/ask_feedback.py (3)
wren-ai-service/src/pipelines/generation/sql_diagnosis.py (1)
run(134-152)wren-ai-service/src/pipelines/generation/sql_correction.py (1)
run(176-206)wren-ai-service/src/web/v1/services/ask.py (1)
AskResult(49-52)
wren-ai-service/src/globals.py (1)
wren-ai-service/src/pipelines/generation/sql_diagnosis.py (1)
SQLDiagnosis(112-152)
wren-ai-service/src/pipelines/generation/sql_diagnosis.py (4)
wren-ai-service/src/core/pipeline.py (1)
BasicPipeline(14-20)wren-ai-service/src/core/provider.py (1)
LLMProvider(6-18)wren-ai-service/src/pipelines/common.py (1)
clean_up_new_lines(111-112)wren-ai-service/src/utils.py (1)
trace_cost(164-182)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
🔇 Additional comments (17)
wren-ai-service/docs/config_examples/config.lm_studio.yaml (1)
146-147: sql_diagnosis added in the right spot; deps look correctSQLDiagnosis is LLM-only, so adding it after sql_tables_extraction with just llm is appropriate. No engine/document_store needed here.
wren-ai-service/docs/config_examples/config.bedrock.yaml (1)
159-160: Bedrock config: sql_diagnosis placement and wiring look goodConsistent with other examples and the pipeline’s LLM-only requirement. Nothing else to change here.
deployment/kustomizations/base/cm.yaml (1)
214-215: Verified end-to-end wiring forsql_diagnosis: class, export, service integration, config examples, and tooling configs are all correctly updated.wren-ai-service/src/pipelines/generation/__init__.py (1)
13-13: Public export for SQLDiagnosis is correctImport + all entry are in line with other pipelines; no circularity concerns from the new module.
Also applies to: 32-32
wren-ai-service/docs/config_examples/config.grok.yaml (1)
148-149: Grok config: sql_diagnosis added correctlyMatches the intended pattern (after sql_tables_extraction; LLM-only). Looks good.
wren-ai-service/tools/config/config.example.yaml (1)
175-176: Approve sql_diagnosis pipe addition — verified wiring and configuration
SQLDiagnosis class is defined and exported, globals registerpipe_components["sql_diagnosis"], and all config examples include the new pipe.wren-ai-service/docs/config_examples/config.google_ai_studio.yaml (1)
152-153: Config example updated with sql_diagnosis — OKThe new pipe entry is aligned with other configs and uses the default LLM as expected.
wren-ai-service/docs/config_examples/config.google_vertexai.yaml (1)
160-161: Vertex AI config: sql_diagnosis entry is consistentGood placement and binding to litellm_llm.default.
wren-ai-service/docs/config_examples/config.deepseek.yaml (1)
166-167: DeepSeek config: sql_diagnosis added — consistentMatches the pattern used across configs; no issues spotted.
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
168-171: ADD_QUOTES path is consistentHere original_sql equals the raw generation_result, which matches the intended semantics.
wren-ai-service/src/globals.py (1)
82-84: Shared SQLDiagnosis instance: confirm statelessness under concurrency.You reuse one _sql_diagnosis_pipeline across services. If any internal component keeps mutable state between runs, this can cause cross-request leakage. Confirm the pipeline is stateless/thread-safe; otherwise instantiate per service.
wren-ai-service/src/web/v1/services/ask.py (2)
433-441: Optional: cachesql_functionsacross retries.You already retrieve
sql_functionsonce; ensure it’s reused inside the retry loop without re-fetching (current code does this—LGTM).
518-548: Great integration of diagnosis-before-correction.The diagnose-then-correct flow with
is_sql_syntax_issuegating reduces futile corrections and should improve latency and quality.wren-ai-service/src/web/v1/services/ask_feedback.py (1)
229-245: LGTM on conditional correction.Only correcting when
is_sql_syntax_issueis true is a sensible scope cut for Ask Feedback.wren-ai-service/src/pipelines/generation/sql_correction.py (1)
60-66: Nice propagation of sql_functions through the prompt and pipeline inputs.Conditional rendering in the template + passing through the prompt and run inputs looks correct and keeps the feature optional.
Also applies to: 89-99, 196-201
wren-ai-service/src/pipelines/generation/sql_diagnosis.py (2)
96-109: Structured output via JSON schema: good.Using a Pydantic model to drive
response_formatand wiring it through the generator is solid and future-proof.Also applies to: 112-131
134-152: No changes required
Downstream callers usesql_diagnosis_results["post_process"].get(...), confirming they expect and handle a Python dict frompost_process.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
wren-ai-service/src/config.py (1)
39-46: Pass settings.allow_sql_diagnosis into AskService in globals.py
In wren-ai-service/src/globals.py’s AskService instantiation (around line 115), addallow_sql_diagnosis=settings.allow_sql_diagnosisso it doesn’t fall back to the service’sTruedefault.wren-ai-service/src/web/v1/services/ask.py (1)
569-572: Avoid None crash on next retry iteration.If correction returns no invalid_generation_result, the next loop iteration will subscript None.
- failed_dry_run_result = sql_correction_results["post_process"][ - "invalid_generation_result" - ] + failed_dry_run_result = sql_correction_results["post_process"][ + "invalid_generation_result" + ] + if not failed_dry_run_result: + break
♻️ Duplicate comments (4)
wren-ai-service/src/web/v1/services/ask.py (2)
505-507: Guard against missing keys in failed_dry_run_result.Direct indexing can KeyError; fall back safely.
- original_sql = failed_dry_run_result["original_sql"] - invalid_sql = failed_dry_run_result["sql"] - error_message = failed_dry_run_result["error"] + original_sql = failed_dry_run_result.get("original_sql") or failed_dry_run_result.get("sql", "") + invalid_sql = failed_dry_run_result.get("sql", "") + error_message = failed_dry_run_result.get("error")
544-549: Ensure non-empty sql and error in correction payload.sql_diagnosis_reasoning can be empty/None; passing None where str is expected may break the correction pipeline.
- invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, - }, + invalid_generation_result={ + "sql": original_sql or invalid_sql or "", + "error": (sql_diagnosis_reasoning || error_message || "SQL diagnosis provided no reasoning") + if allow_sql_diagnosis + else (error_message || "Unknown SQL error"), + },wren-ai-service/src/web/v1/services/ask_feedback.py (2)
194-197: Prevent KeyError on missing fields from invalid_generation_result.- original_sql = failed_dry_run_result["original_sql"] - invalid_sql = failed_dry_run_result["sql"] - error_message = failed_dry_run_result["error"] + original_sql = failed_dry_run_result.get("original_sql") or failed_dry_run_result.get("sql", "") + invalid_sql = failed_dry_run_result.get("sql", "") + error_message = failed_dry_run_result.get("error")
222-236: Harden correction payload with safe fallbacks.- invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, - }, + invalid_generation_result={ + "sql": original_sql or invalid_sql or "", + "error": (sql_diagnosis_reasoning || error_message || "SQL diagnosis provided no reasoning") + if allow_sql_diagnosis + else (error_message || "Unknown SQL error"), + },
🧹 Nitpick comments (3)
wren-ai-service/src/web/v1/services/ask.py (1)
97-106: Constructor default should match Settings.default (False).To keep behavior consistent across entry points, make allow_sql_diagnosis default False (or always pass it from Settings).
- allow_sql_diagnosis: bool = True, + allow_sql_diagnosis: bool = False,wren-ai-service/src/web/v1/services/ask_feedback.py (2)
59-66: Constructor default should match Settings.default (False).Keep allow_sql_diagnosis default consistent across Settings/Ask/AskFeedback.
- allow_sql_diagnosis: bool = True, + allow_sql_diagnosis: bool = False,
238-253: Optional: use .get() when re-reading failed correction result.Minor resilience; avoids subscripting errors if schema changes.
- invalid_sql = failed_dry_run_result["sql"] - error_message = failed_dry_run_result["error"] + invalid_sql = failed_dry_run_result.get("sql", "") + error_message = failed_dry_run_result.get("error")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(3 hunks)wren-ai-service/src/web/v1/services/ask.py(5 hunks)wren-ai-service/src/web/v1/services/ask_feedback.py(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/src/globals.py
🧰 Additional context used
🧬 Code graph analysis (2)
wren-ai-service/src/web/v1/services/ask_feedback.py (4)
wren-ai-service/src/pipelines/generation/sql_diagnosis.py (1)
run(134-152)wren-ai-service/src/pipelines/generation/sql_correction.py (1)
run(176-206)wren-ai-service/src/pipelines/retrieval/sql_functions.py (1)
run(104-125)wren-ai-service/src/web/v1/services/ask.py (1)
AskResult(49-52)
wren-ai-service/src/web/v1/services/ask.py (1)
wren-ai-service/src/pipelines/generation/sql_correction.py (1)
run(176-206)
🔇 Additional comments (8)
wren-ai-service/src/web/v1/services/ask.py (4)
115-120: Good: internal flag stored once.The private flag mirrors the constructor parameter cleanly.
168-172: Local snapshot of flags is fine.Capturing allow_sql_diagnosis for flow control improves readability.
521-538: Diagnosis handling: clarify None/unknown semantics.If is_sql_syntax_issue is None (missing), the current logic breaks the loop. Confirm that “unknown” should be treated as non-syntax to skip correction; otherwise default to attempting correction.
553-553: Nice: pass sql_functions into correction.This keeps corrections aware of available functions.
wren-ai-service/src/web/v1/services/ask_feedback.py (4)
71-73: Good: store flag.
100-102: Local snapshots are fine.
133-141: Good: function retrieval gated by flag.
205-221: Diagnosis default semantics.You default is_sql_syntax_issue to True, then overwrite. If the key is absent (None), correction is skipped. Confirm this is intended “fail-closed” behavior.
| invalid_sql = failed_dry_run_result["sql"] | ||
| error_message = failed_dry_run_result["error"] |
There was a problem hiding this comment.
🛠️ Refactor suggestion
TIME_OUT branch: also guard against missing fields.
- invalid_sql = failed_dry_run_result["sql"]
- error_message = failed_dry_run_result["error"]
+ invalid_sql = failed_dry_run_result.get("sql", "")
+ error_message = failed_dry_run_result.get("error")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| invalid_sql = failed_dry_run_result["sql"] | |
| error_message = failed_dry_run_result["error"] | |
| invalid_sql = failed_dry_run_result.get("sql", "") | |
| error_message = failed_dry_run_result.get("error") |
🤖 Prompt for AI Agents
In wren-ai-service/src/web/v1/services/ask_feedback.py around lines 255 to 256,
the TIME_OUT branch assumes failed_dry_run_result contains "sql" and "error"
keys and can raise KeyError; change the code to safely access these fields using
dictionary .get(...) with sensible defaults (e.g., empty string or a descriptive
fallback) and/or check for key presence before using them so invalid_sql and
error_message are always defined and non-None when used further.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (5)
wren-ai-service/src/pipelines/generation/utils/sql.py (2)
127-129: Preserve raw original SQL for diagnosis (use generation_result, not quoted_sql).original_sql should capture the unquoted LLM output so downstream diagnosis works on the true source text.
- "sql": addition.get("error_sql", quoted_sql), - "original_sql": quoted_sql, + "sql": addition.get("error_sql", quoted_sql), + "original_sql": generation_result,
157-159: Same: keep unmodified original SQL in PREVIEW error path.Mirror the change here so preview errors also retain the raw input.
- "sql": addition.get("error_sql", quoted_sql), - "original_sql": quoted_sql, + "sql": addition.get("error_sql", quoted_sql), + "original_sql": generation_result,wren-ai-service/src/pipelines/generation/sql_diagnosis.py (1)
93-98: post_process: fix return type and harden parsing with schema validation.
- Annotated as -> str but returns a parsed object.
- Accessing replies[0] without checks may raise.
Use the Pydantic schema to validate.-@observe(capture_input=False) -async def post_process( - generate_sql_diagnosis: dict, -) -> str: - return orjson.loads(generate_sql_diagnosis.get("replies")[0]) +@observe(capture_input=False) +async def post_process( + generate_sql_diagnosis: dict, +) -> dict: + replies = generate_sql_diagnosis.get("replies") or [] + if not replies: + raise ValueError("SQLDiagnosis: generator returned no replies") + try: + return SqlDiagnosisResult.model_validate_json(replies[0]).model_dump() + except Exception as e: + logger.exception("SQLDiagnosis: failed to parse JSON reply: %s", e) + raisewren-ai-service/src/web/v1/services/ask.py (2)
505-507: Guard against missing keys in failed_dry_run_result.Upstream may omit original_sql; avoid KeyError and ensure sensible fallbacks.
- original_sql = failed_dry_run_result["original_sql"] - invalid_sql = failed_dry_run_result["sql"] - error_message = failed_dry_run_result["error"] + original_sql = failed_dry_run_result.get("original_sql") or failed_dry_run_result.get("sql", "") + invalid_sql = failed_dry_run_result.get("sql", "") + error_message = failed_dry_run_result.get("error")
545-550: Harden correction payload: ensure non-empty sql and error message.Fallback to invalid_sql and original error when diagnosis is disabled or empty.
- invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, - }, + invalid_generation_result={ + "sql": original_sql or invalid_sql, + "error": (sql_diagnosis_reasoning or error_message) if allow_sql_diagnosis else (error_message or "SQL diagnosis unavailable"), + },
🧹 Nitpick comments (5)
wren-ai-service/src/pipelines/generation/utils/sql.py (3)
71-79: Fix return type annotation of _classify_generation_result.The function returns a 2-tuple, but the annotation says Dict[str, str]. Update to a precise tuple type to avoid misleading callers and static analysis tools.
- ) -> Dict[str, str]: + ) -> tuple[dict, dict]:
41-46: Be robust when detecting JSON-wrapped replies.Relying on startswith("{") can misclassify content with leading whitespace. Prefer a try/except parse.
- if cleaned_generation_result.startswith("{"): - cleaned_generation_result = orjson.loads(cleaned_generation_result)[ - "sql" - ] + try: + parsed = orjson.loads(cleaned_generation_result) + if isinstance(parsed, dict) and "sql" in parsed: + cleaned_generation_result = parsed["sql"] + except orjson.JSONDecodeError: + pass
233-235: Nit: fix typo in RULES string.“table/columnalias” → “table/column alias”.
- - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. + - ONLY USE table/column alias in the final SELECT clause; don't use table/column alias in the other clauses.wren-ai-service/src/pipelines/generation/sql_diagnosis.py (1)
85-91: Return type annotation does not match actual return.Function returns (result, generator_name) for trace_cost, but annotation says dict. Align the signature.
-async def generate_sql_diagnosis( - prompt: dict, generator: Any, generator_name: str -) -> dict: +async def generate_sql_diagnosis( + prompt: dict, generator: Any, generator_name: str +) -> tuple[dict, str]:wren-ai-service/src/web/v1/services/ask.py (1)
521-539: Optional: guard diagnosis outputs.If the diagnosis pipeline returns an empty object, .get(...) is fine; consider defaulting can_be_corrected to True to avoid prematurely breaking out. Not blocking.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (6)
wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_correction.py(5 hunks)wren-ai-service/src/pipelines/generation/sql_diagnosis.py(1 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(3 hunks)wren-ai-service/src/web/v1/services/ask.py(5 hunks)wren-ai-service/src/web/v1/services/ask_feedback.py(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- wren-ai-service/src/config.py
- wren-ai-service/src/pipelines/generation/sql_correction.py
- wren-ai-service/src/web/v1/services/ask_feedback.py
🧰 Additional context used
🧬 Code graph analysis (3)
wren-ai-service/src/web/v1/services/ask.py (2)
wren-ai-service/src/pipelines/generation/sql_correction.py (1)
run(157-187)wren-ai-service/src/pipelines/generation/utils/sql.py (1)
run(29-69)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
wren-ai-service/src/providers/engine/wren.py (3)
execute_sql(26-123)execute_sql(143-191)execute_sql(262-312)
wren-ai-service/src/pipelines/generation/sql_diagnosis.py (5)
wren-ai-service/src/core/pipeline.py (1)
BasicPipeline(14-20)wren-ai-service/src/core/provider.py (1)
LLMProvider(6-18)wren-ai-service/src/pipelines/common.py (1)
clean_up_new_lines(111-112)wren-ai-service/src/utils.py (1)
trace_cost(164-182)wren-ai-service/src/pipelines/generation/sql_correction.py (3)
prompt(81-96)run(157-187)post_process(108-122)
🔇 Additional comments (2)
wren-ai-service/src/web/v1/services/ask.py (2)
104-119: New allow_sql_diagnosis flag wiring looks good.Constructor and instance wiring are consistent with existing flags.
170-171: Localizing allow_sql_diagnosis looks good.Using the instance-level flag is consistent with other toggles.
2 participants
Summary by CodeRabbit
New Features
Improvements
Documentation
Chores