chore(wren-ai-service): allow regenerate sql using retrieved tables #1324
Conversation
cyyeh
added
module/ai-service
WalkthroughThis set of changes updates several core modules by modifying function signatures to include lists of tables, integrating database schema documents into SQL generation prompts, and simplifying control flows by removing redundant user correction functions. The retrieval pipelines now handle empty queries better and support filtering, while the document store adds a new asynchronous query method. Additionally, status values in feedback responses have been updated to better reflect processing stages. Changes
Sequence Diagram(s)sequenceDiagram
participant U as User
participant UI as UI (on_click_regenerate_sql)
participant S as AskFeedback Service
participant R as Retrieval Pipeline
participant Q as Qdrant Document Store
participant G as SQL Regeneration Pipeline
U->>UI: Provide comma-separated table names & click regenerate
UI->>S: Call ask_feedback(tables, sql_generation_reasoning, sql)
S->>R: Invoke Retrieval with tables and SQL reasoning
R->>Q: Query documents using embedding or filters
Q-->>R: Return retrieved documents
R->>G: Pass retrieved table DDLs as contexts
G-->>R: Return regenerated SQL query
R-->>S: Return feedback results
S->>UI: Send AskFeedbackResultResponse ("searching")
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: 1
🔭 Outside diff range comments (2)
wren-ai-service/src/pipelines/retrieval/retrieval.py (1)
138-165: 🛠️ Refactor suggestionConsider validating tables parameter.
The table_retrieval function should validate the tables parameter before using it in filters.
async def table_retrieval( embedding: dict, id: str, tables: list[str], table_retriever: Any ) -> dict: + if tables is not None and not isinstance(tables, list): + raise ValueError("tables parameter must be a list of strings") + if tables is not None and not all(isinstance(t, str) for t in tables): + raise ValueError("all table names must be strings") filters = { "operator": "AND", "conditions": [wren-ai-service/src/web/v1/services/ask.py (1)
642-642:⚠️ Potential issueFix: SQL correction uses empty contexts.
The SQL correction pipeline is called with empty contexts (
contexts=[]), which could lead to incorrect SQL corrections since the table DDLs are not passed through.Apply this fix:
- contexts=[], + contexts=table_ddls,
🧹 Nitpick comments (1)
wren-ai-service/demo/utils.py (1)
232-236: Consider making the table input handling more robust.While the comma-separated input is user-friendly, the current splitting logic might be sensitive to extra spaces or different separators.
Consider this more robust implementation:
- retrieved_tables.split(", "), + [table.strip() for table in retrieved_tables.split(",") if table.strip()],Also applies to: 250-253
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
wren-ai-service/demo/utils.py(5 hunks)wren-ai-service/src/pipelines/generation/sql_regeneration.py(6 hunks)wren-ai-service/src/pipelines/retrieval/retrieval.py(5 hunks)wren-ai-service/src/providers/document_store/qdrant.py(3 hunks)wren-ai-service/src/web/v1/routers/ask.py(1 hunks)wren-ai-service/src/web/v1/services/ask.py(4 hunks)
⏰ 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 (14)
wren-ai-service/src/pipelines/generation/sql_regeneration.py (4)
26-26: LGTM! Enhanced system prompt with database schema context.The system prompt now correctly instructs the model to utilize the database schema when generating SQL queries.
Also applies to: 29-29
42-45: LGTM! Added database schema section to user prompt template.The template now properly includes database schema information through document iteration.
63-63: LGTM! Updated prompt function signature with documents parameter.The function signature and implementation correctly handle the new documents parameter for database schema context.
Also applies to: 72-74
140-142: Verify error handling for empty contexts list.The run method now accepts contexts but should handle cases where the contexts list is empty.
# Add error handling for empty contexts async def run( self, contexts: list[str], sql_generation_reasoning: str, sql: str, ... ): if not contexts: logger.warning("No database schema contexts provided for SQL regeneration")Also applies to: 152-165
wren-ai-service/src/web/v1/routers/ask.py (1)
158-158: LGTM! Updated status to accurately reflect the operation state.Changed status from "understanding" to "searching" to better represent the current operation phase.
wren-ai-service/src/providers/document_store/qdrant.py (1)
350-362: LGTM! Enhanced query logic with fallback to filters.The run method now intelligently chooses between embedding-based and filter-based queries.
wren-ai-service/src/pipelines/retrieval/retrieval.py (3)
121-133: LGTM! Added proper handling for empty queries.The function now correctly handles empty queries by returning an empty dictionary.
170-171: LGTM! Simplified dbschema_retrieval signature.Removed unused embedding parameter and simplified the function signature.
Also applies to: 196-196
482-484: LGTM! Updated run method with optional parameters.The run method now properly handles optional query and tables parameters.
Also applies to: 492-492
wren-ai-service/src/web/v1/services/ask.py (3)
105-105: LGTM! Well-structured addition of tables attribute.The new
tablesattribute is properly typed and logically placed within the request model.
149-149: LGTM! Status update accurately reflects the operation.The status change from "understanding" to "searching" better represents the actual operation being performed during feedback processing.
593-601: LGTM! Retrieval logic properly updated for table-specific search.The retrieval pipeline is correctly modified to use the provided tables, and table DDLs are properly extracted from the results.
wren-ai-service/demo/utils.py (2)
186-188: LGTM! Function signature properly updated.The addition of the
retrieved_tablesparameter is well-typed and correctly implements the new functionality.
533-533: LGTM! Function properly updated to handle table lists.The
ask_feedbackfunction is correctly modified to accept and pass through the table list parameter.Also applies to: 537-537
| while True: | ||
| points = await self.async_client.scroll( | ||
| collection_name=self.index, | ||
| offset=offset, | ||
| scroll_filter=qdrant_filters, | ||
| limit=top_k, | ||
| ) | ||
| points_list.extend(points[0]) | ||
| if points[1] is None: | ||
| break | ||
| offset = points[1] |
There was a problem hiding this comment.
I executed the code and found that None is present at index 1, allowing it to satisfy the condition and exit the while loop. However, I suggest we clarify this in the docstring for future reference.
There was a problem hiding this comment.
Additionally, the while loop w/ an always True condition poses a potential risk of an infinite loop. we are better to change the condition from my opinion.
POST /v1/ask-feedbacks: now users need to pass parametertables(type: [string], which means table names) to regenerate sql.Summary by CodeRabbit
New Features
Improvements