chore(wren-ai-service): add retrieved_tables to sql correction api #1804
Conversation
cyyeh
added
module/ai-service
WalkthroughThis change introduces an optional Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant API
participant SqlCorrectionService
participant Pipeline
Client->>API: POST /sql_corrections (with optional retrieved_tables)
API->>SqlCorrectionService: correct(request)
alt retrieved_tables provided
SqlCorrectionService->>SqlCorrectionService: Use request.retrieved_tables
else not provided
SqlCorrectionService->>Pipeline: Run "sql_tables_extraction"
Pipeline-->>SqlCorrectionService: Extracted tables
end
SqlCorrectionService->>Pipeline: Run "db_schema_retrieval" with tables
Pipeline-->>SqlCorrectionService: Schema
SqlCorrectionService-->>API: Correction result
API-->>Client: Response
Possibly related PRs
Suggested labels
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. 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 (2)
wren-ai-service/src/web/v1/routers/sql_corrections.py (1)
19-32: Update documentation to include the new field.The router documentation should be updated to reflect the new
retrieved_tablesfield in the request body schema.{ "sql": "SELECT * FROM table", # Invalid SQL statement "error": "Error message" # Error message - "project_id": "project-id" # Optional project ID + "project_id": "project-id", # Optional project ID + "retrieved_tables": ["table1", "table2"] # Optional list of table names }wren-ai-service/src/web/v1/services/sql_corrections.py (1)
88-104: Consider error handling for invalid table names.While the logic correctly bypasses extraction when tables are provided, consider adding validation to ensure the provided table names are valid and exist in the database schema. Invalid table names could cause downstream pipeline failures.
Example approach:
if not retrieved_tables: retrieved_tables = ( await self._pipelines["sql_tables_extraction"].run( sql=sql, ) )["post_process"] else: # Optional: Validate table names exist in the database logger.info(f"Using provided tables: {retrieved_tables}")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
wren-ai-service/src/web/v1/routers/sql_corrections.py(2 hunks)wren-ai-service/src/web/v1/services/sql_corrections.py(4 hunks)
🧰 Additional context used
🧠 Learnings (2)
wren-ai-service/src/web/v1/routers/sql_corrections.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/sql_corrections.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.
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
🔇 Additional comments (6)
wren-ai-service/src/web/v1/routers/sql_corrections.py (2)
3-3: LGTM: Import addition is correct.The
Listimport is properly added to support the newretrieved_tablesfield type.
70-70: LGTM: Field addition follows established patterns.The
retrieved_tablesfield is correctly defined as optional with appropriate typing.wren-ai-service/src/web/v1/services/sql_corrections.py (4)
2-2: LGTM: Import addition is correct.The
Listimport is properly added to support the newretrieved_tablesfield type.
61-61: LGTM: Field addition follows established patterns.The
retrieved_tablesfield is correctly defined as optional with appropriate typing, consistent with the router model.
78-78: LGTM: Variable assignment is correct.The
retrieved_tablesvariable is properly assigned from the request.
99-99: LGTM: Pipeline parameter updated correctly.The parameter is correctly updated to use
retrieved_tablesinstead of the previously extractedtablesvariable.
Summary by CodeRabbit