feat(wren-ai-service): Add SQL Correction Service and API Endpoints#1420
Conversation
paopa
added
module/ai-service
WalkthroughThis changeset introduces a new SQL correction service into the application. A new Changes
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
✨ Finishing Touches
🪧 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 (4)
wren-ai-service/src/web/v1/routers/sql_corrections.py (2)
76-97: Ensure pipeline availability and consider concurrency.
- The code references
service_container.sql_correction_serviceand background tasks without verifying that thesql_correctionpipeline exists. If it’s misconfigured or missing, aKeyErrormight occur.- Consider whether unlimited concurrent SQL correction tasks might strain resources. You could implement a queue or concurrency limit to prevent potential load spikes.
Would you like me to provide a script to search for pipeline instantiation references in the codebase to ensure that
"sql_correction"is always defined?
107-113: Consider returning a 404 for unknown events.Currently, if the
event_idis unknown, the service returns a default failed status. From an API design perspective, returning HTTP 404 might be more intuitive to indicate that the event does not exist.wren-ai-service/src/web/v1/services/sql_corrections.py (2)
27-35: Reevaluate cache size and TTL.
maxsize=1_000_000andttl=120can quickly consume memory if there are frequent requests with minimal cache rotation. Consider adjusting the limits or setting up monitoring to ensure stability under high load.
36-50: Consider enriching log details on exceptions.
_handle_exceptioncurrently logs only the message. It might be helpful to include more details (e.g., stack traces) to accelerate troubleshooting in production environments.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
wren-ai-service/src/globals.py(2 hunks)wren-ai-service/src/web/v1/routers/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/sql_corrections.py(1 hunks)wren-ai-service/src/web/v1/services/__init__.py(2 hunks)wren-ai-service/src/web/v1/services/instructions.py(2 hunks)wren-ai-service/src/web/v1/services/sql_corrections.py(1 hunks)
🧰 Additional context used
🧬 Code Definitions (2)
wren-ai-service/src/globals.py (3)
wren-ai-service/src/pipelines/generation/sql_correction.py (1) (1)
SQLCorrection(99:143)
wren-ai-service/src/pipelines/generation/sql_correction.py (1) (1)
SQLCorrection(99:143)
wren-ai-service/src/pipelines/generation/sql_correction.py (1) (1)
SQLCorrection(99:143)
wren-ai-service/src/web/v1/services/sql_corrections.py (2)
wren-ai-service/src/web/v1/routers/sql_corrections.py (2) (2)
correct(77:96)get(108:113)
wren-ai-service/src/web/v1/routers/sql_corrections.py (2) (2)
correct(77:96)get(108:113)
⏰ Context from checks skipped due to timeout of 90000ms (4)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (12)
wren-ai-service/src/web/v1/services/__init__.py (2)
72-72: Implementation looks goodThe import for the new SqlCorrectionService is properly added, following the established pattern in the codebase.
87-87: Correctly exported the new serviceThe SqlCorrectionService is appropriately added to the all list, making it available for import by other modules.
wren-ai-service/src/web/v1/routers/__init__.py (2)
14-14: Import looks goodThe sql_corrections module is properly imported, consistent with the pattern for other routers.
34-34: Router integration is correctThe sql_corrections router is properly integrated into the main router, maintaining the established pattern in the codebase.
wren-ai-service/src/web/v1/services/instructions.py (2)
54-54: Error instantiation refactoring looks goodChanged from
self.Event.Errortoself.Error, simplifying the error instantiation approach and likely aligning with patterns used in the new SqlCorrectionService.
154-154: Consistent error handling approachUpdated to use
self.Errordirectly, matching the change made in the _handle_exception method and maintaining consistency throughout the service.wren-ai-service/src/globals.py (2)
30-30: Service container attribute looks goodThe SqlCorrectionService is properly added to the ServiceContainer class, maintaining consistent structure with other services.
261-269: Service instantiation looks goodThe SqlCorrectionService is properly instantiated with the correct pipeline configuration, following the established pattern in the codebase. The service reuses the existing sql_correction pipeline that's already used elsewhere in the application (e.g., in sql_expansion_service), which is a good practice for code reuse.
wren-ai-service/src/web/v1/routers/sql_corrections.py (2)
19-63: Good documentation for the new endpoints.The detailed docstring provides clarity regarding the endpoint usage and workflow, which is helpful for both developers and API consumers.
66-70: Validate optional fields.
project_idis optional, which is fine. However, consider validating or normalizing the project ID if it's a critical piece of metadata. This ensures that downstream services won't fail due to an unexpected format or missing value.wren-ai-service/src/web/v1/services/sql_corrections.py (2)
15-25: Clear event modeling.The
ErrorandEventmodels neatly encapsulate all necessary properties, facilitating straightforward data handling. This approach makes it easy to track the status of each correction request and attach relevant error information.
107-120: Consistent approach for unknown events.The fallback to a failed event is consistent with your design. However, if you decide to return a 404 in the router, ensure that the same logic is reflected here by throwing a custom exception or returning a sentinel value that the router can interpret.
paopa
force-pushed
the
feat/sql-correcting-endpoint
branch
from
383d804 to
caa3f87
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (9)
wren-ai-service/src/globals.py (2)
30-30: Add a concise docstring or comment for “sql_correction_service.”
Currently, there is no description for this attribute. Including a docstring or comment clarifying its purpose would improve maintainability and readability.
261-269: Consider consolidating repeated pipeline instances.
Asql_correctionpipeline is added here, but a similar pipeline reference appears in other services (e.g.,ask_serviceat line 116). Reusing or consolidating this pipeline might avoid duplication and reduce overhead.wren-ai-service/src/web/v1/routers/sql_corrections.py (2)
76-97: Background task usage is correct; consider adding error response.
While the primary logic is in the background task, you may optionally provide an immediate check that the pipeline key exists (like you do in the service). If missing, you might return a 400 or informative response. Currently, the code always defers checking to the background task.
107-113: Return 404 or similar for missing events.
Currently, if an event is missing, the code returns afailedstatus withOTHERSerror. Consider using HTTP status 404 to reflect that the resource (event) does not exist.wren-ai-service/src/web/v1/services/sql_corrections.py (5)
16-19: Error model is sufficient for general failures.
Provides a helpful structure (code,message). If you anticipate more error codes, consider enumerating them.
27-35: Cache-based storage is straightforward but be mindful of concurrency.
TTLCacheis not inherently threadsafe if used in a multi-worker environment. If concurrency grows, you may need a lock or a thread-safe cache variant.
36-50: Exception handling logs the error but consider advanced logging strategies.
When an exception occurs, you set the event to “failed” and store the error. This is good. For large-scale debugging, consider structured logging or linking a correlation ID.
57-106: Check for missing pipeline keys before usage.
You doself._pipelines["sql_correction"]directly. Ifsql_correctionis unexpectedly absent, you raise a KeyError. You have a fallback in_handle_exception, but a pre-check with a friendlier error might improve resilience.
107-120: Graceful handling of missing or expired events.
Returns a “failed” event with message. This is user-friendly, but consider also returning an HTTP 404 in the router if an event is absent.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
wren-ai-service/src/globals.py(2 hunks)wren-ai-service/src/web/v1/routers/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/sql_corrections.py(1 hunks)wren-ai-service/src/web/v1/services/__init__.py(2 hunks)wren-ai-service/src/web/v1/services/instructions.py(2 hunks)wren-ai-service/src/web/v1/services/sql_corrections.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- wren-ai-service/src/web/v1/services/instructions.py
- wren-ai-service/src/web/v1/services/init.py
- wren-ai-service/src/web/v1/routers/init.py
🧰 Additional context used
🧬 Code Definitions (3)
wren-ai-service/src/web/v1/routers/sql_corrections.py (1)
wren-ai-service/src/web/v1/services/sql_corrections.py (4) (4)
SqlCorrectionService(15:122)correct(59:105)Event(20:25)CorrectionRequest(51:55)
wren-ai-service/src/globals.py (1)
wren-ai-service/src/pipelines/generation/sql_correction.py (1) (1)
SQLCorrection(99:143)
wren-ai-service/src/web/v1/services/sql_corrections.py (1)
wren-ai-service/src/web/v1/routers/sql_corrections.py (2) (2)
correct(77:96)get(108:113)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (9)
wren-ai-service/src/web/v1/routers/sql_corrections.py (5)
1-14: Imports and module setup look good.
All necessary imports for FastAPI, Pydantic, and service dependencies are properly organized. No issues noted here.
16-63: Docstring thoroughly documents the router’s purpose.
The router-level docstring is comprehensive and outlines usage, endpoints, and request/response models. Nicely done for maintainability.
66-70: Validate nullable fields in PostRequest.
project_idis optional, but consider confirming that an empty string vs. null is handled appropriately downstream.
72-74: PostResponse data structure is minimal but sufficient.
Onlyevent_idis returned. This is clear for asynchronous processes. No issues found.
99-105: GetResponse model is straightforward.
It closely parallels the event’s structure in SqlCorrectionService. No issues found.wren-ai-service/src/web/v1/services/sql_corrections.py (4)
1-4: Imports and logging setup are fine.
No extraneous imports or unused references detected here.
20-26: Event model design is clear.
Covers typical states for asynchronous tasks, including an optionaltrace_idfor debugging. Good approach.
51-56: CorrectionRequest structure is well-defined.
Captures essentials (sql,error,project_id). This is adequate for passing correction data.
121-123: setitem usage is consistent with your caching approach.
Directly assigning theEventobject is fine; just keep concurrency considerations in mind if usage increases.
cyyeh
left a comment
cyyeh
left a comment
There was a problem hiding this comment.
TODO: we'll improve this PR by adding sql functions, instructions etc.
2 participants
This PR introduces a new SQL Correction Service that helps users fix invalid SQL queries through a dedicated API endpoint.
Key changes:
SqlCorrectionServicewith background task processing/sql-correctionsendpoints for POST and GET operationsThe service accepts invalid SQL queries with error messages and returns corrected SQL statements through an asynchronous process. Users can track the correction status using the provided event ID.
Endpoints:
/sql-corrections: Submit invalid SQL for correction/sql-corrections/{event_id}: Check correction status and resultsSummary by CodeRabbit