chore(wren-ai-service): improve sql pairs and instructions#1422
Conversation
WalkthroughThe changes update multiple modules in the AI service pipeline. In the generation pipelines, the Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant AskService
participant SQLPipeline as "SQL Samples Pipeline"
participant InstructionPipeline as "Instructions Pipeline"
participant IntentClassification
Client->>AskService: Send query request
AskService->>SQLPipeline: Retrieve SQL samples (async)
AskService->>InstructionPipeline: Retrieve instructions (async)
SQLPipeline-->>AskService: Return SQL samples
InstructionPipeline-->>AskService: Return instructions
AskService->>IntentClassification: Process query with samples & instructions
IntentClassification-->>AskService: Return classification result
AskService-->>Client: Return aggregated response
Possibly related PRs
Poem
Tip ⚡🧪 Multi-step agentic review comment chat (experimental)
✨ 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: 2
🧹 Nitpick comments (9)
wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py (2)
130-132: Consider extracting the queue creation logic into a helper.These lines are duplicated in
_streaming_callbackandget_streaming_results. A small helper function (e.g.,_ensure_user_queue) would improve maintainability by eliminating redundant code.
143-145: Duplicate queue creation logic here as well.Similar to lines 130-132, this block reimplements the queue creation. Consider the same helper approach to keep the code DRY and easier to maintain.
wren-ai-service/src/web/v1/services/ask.py (4)
214-214: Avoid hardcoding an empty list if instructions are not given.You might default to an empty list in function parameters or handle it when passing arguments downstream. This ensures clarity about how "instructions" are set and reduces confusion in future refactoring.
267-272: Double-check assumptions about task outputs.Here you map
sql_samples_task["formatted_output"]andinstructions_task["formatted_output"]to local variables. If the tasks return unexpected keys or structures, this may raise KeyError or TypeError. Consider validating the data shape before using it.
274-275: Consider logging or tracing the intent classification parameter.When
self._allow_intent_classificationis True, we proceed to intent classification. Logging a message here (e.g. “intent classification started”) could help debug production issues if classification is unexpectedly skipped or triggered.
405-405: Check compatibility of instructions in sql_generation.Similar to line 394, you’re passing
instructionsintosql_generation. Confirm that the pipeline handles or validates these instructions (e.g., logs them, ensures they are not empty).wren-ai-service/src/pipelines/generation/intent_classification.py (1)
114-119: Document or validate the nature of instructions.The template loops over each
instructionand prints it. Make sure these instructions are consistently typed as strings and that no further structure is expected.wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py (2)
53-58: New instructions section added to templateThe addition of a conditional instructions section enhances the template's flexibility. This allows for dynamic behavior based on whether instructions are provided.
Consider adding code comments or documentation about what kind of instructions are expected and how they will be used in the reasoning process. This would help future developers understand the expected format and purpose of these instructions.
201-210: Consider updating dry_run_pipeline callThe dry_run_pipeline call doesn't include the new instructions parameter. While this won't cause issues since it's optional, consider updating it for consistency and testing purposes.
dry_run_pipeline( FollowUpSQLGenerationReasoning, "followup_sql_generation_reasoning", query="this is a test query", histories=[], contexts=[], + instructions=["Sample instruction"], )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py(5 hunks)wren-ai-service/src/pipelines/generation/intent_classification.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py(2 hunks)wren-ai-service/src/web/v1/services/ask.py(4 hunks)wren-ai-service/src/web/v1/services/sql_explanation.py(0 hunks)wren-ai-service/src/web/v1/services/sql_regeneration.py(0 hunks)
💤 Files with no reviewable changes (2)
- wren-ai-service/src/web/v1/services/sql_explanation.py
- wren-ai-service/src/web/v1/services/sql_regeneration.py
⏰ 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 (9)
wren-ai-service/src/web/v1/services/ask.py (3)
253-255: Validate concurrent tasks more robustly.You call
asyncio.gatherto fetchsql_pairs_retrievalandinstructions_retrievalsimultaneously. If one task fails, the entire call may raise an exception, masking partial results. Consider addingreturn_exceptions=Trueor explicit error handling blocks to gracefully degrade if one retrieval fails.Also applies to: 256-263
280-281: Passing new parameters to intent classification.These two lines effectively ensure
sql_samplesandinstructionsare provided to the pipeline. Confirm that downstream usage inintent_classificationmatches the expected format (e.g., list of dict vs list of strings) to avoid runtime confusion.
394-394: Maintain consistency when passing instructions.Providing
instructions=instructionsin the follow-up call is a good step toward consistent usage. However, ensure all callers of this pipeline also handleinstructionsconsistently (e.g., default empty list vs None).wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py (6)
50-51: Proper formatting in SQL samples sectionThe code properly closes the iteration and conditional block for SQL samples, maintaining consistent template structure.
60-66: Improved query history sectionRenaming from "CONTEXT" to "User's QUERY HISTORY" and reformatting with explicit "Question:" and "SQL:" labels improves clarity for both the LLM and developers.
84-84: Added instructions parameter to prompt functionThe function signature has been properly updated to include the new instructions parameter.
93-93: Instructions parameter passed to prompt builderThe instructions parameter is correctly passed to the prompt_builder.run() function, ensuring it's incorporated into the generated prompt.
181-181: Added optional instructions parameter to run methodThe run method signature has been appropriately updated to accept an optional instructions parameter, maintaining backward compatibility.
193-193: Instructions parameter handled in execute inputsThe instructions parameter is correctly handled with a default empty list when None is provided, ensuring the template can process it properly.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
wren-ai-service/src/pipelines/generation/sql_generation.py (1)
63-64:Details
❓ Verification inconclusive
Updated type annotations to use native Python types
Type annotations have been modernized by replacing
List[str]withlist[str]andList[Dict]withlist[dict]. This follows Python 3.9+ best practices for type hinting.However, there's a significant change in the
instructionsparameter type fromList[Dict] | Nonetolist[str] | None, which alters the expected data structure.Also applies to: 67-69
🏁 Script executed:
#!/bin/bash # Check if the construct_instructions function correctly handles the new instructions type grep -r "def construct_instructions" --include="*.py" .Length of output: 141
Verify that the updated type for
instructionsis intentional and compatible across the pipeline.
- The changes replacing
List[str]withlist[str](and similarly for dictionaries) follow Python 3.9+ best practices.- However, the modification of the
instructionsparameter type—fromList[Dict] | Nonetolist[str] | None—is significant. Theconstruct_instructionsfunction (located inwren-ai-service/src/pipelines/generation/utils/sql.py) appears to be downstream of this change.- Please review and confirm that
construct_instructions(and any other related functionality) is correctly updated to handle alist[str] | Noneinput rather than a list of dictionaries.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
wren-ai-service/src/pipelines/generation/followup_sql_generation.py(3 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py(6 hunks)wren-ai-service/src/pipelines/generation/intent_classification.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(3 hunks)wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py(5 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- wren-ai-service/src/pipelines/generation/followup_sql_generation.py
🚧 Files skipped from review as they are similar to previous changes (3)
- wren-ai-service/src/pipelines/generation/intent_classification.py
- wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py
- wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py
🧰 Additional context used
🧬 Code Definitions (1)
wren-ai-service/src/pipelines/generation/sql_generation.py (2)
wren-ai-service/src/web/v1/services/__init__.py (1) (1)
Configuration(25:45)wren-ai-service/src/providers/__init__.py (1) (1)
Configuration(265:267)
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: pytest
- GitHub Check: Analyze (go)
🔇 Additional comments (4)
wren-ai-service/src/pipelines/generation/sql_generation.py (4)
3-3: Modernized import statementsThe imports have been updated to use only
AnyandOptionalfrom thetypingmodule, removingListandDict. This aligns with Python 3.9+ best practices where built-in collection types can be used directly in type annotations.
144-145: Consistent type annotation updates in the run methodThe type annotations in the
runmethod have been updated to match those in thepromptfunction, ensuring consistency throughout the codebase. The same change fromList[Dict] | Nonetolist[str] | Nonefor theinstructionsparameter is applied here.Also applies to: 147-149
158-165: Inputs dictionary correctly includes all parametersThe
runmethod correctly passes all parameters to the pipeline execution, including the newly updatedinstructionsparameter. This ensures the type changes are consistently propagated through the execution flow.
30-33:Details
❓ Verification inconclusive
Instructions template usage matches new type
The template uses
instructionsdirectly without iteration, which aligns with the changed type from a list of dictionaries to a list of strings. However, it's not clear how multiple instructions in the list would be rendered - they would likely be concatenated or joined with newlines.
🏁 Script executed:
#!/bin/bash # Check if there are calls to this function with instructions as list[str] grep -r "instructions=" --include="*.py" . | grep -v "instructions=None" | grep -v "instructions=construct_instructions"Length of output: 2971
Review Note: Verify Multiple Instructions Formatting
The template now directly outputs the
instructionsvariable, which aligns with the change from a list of dictionaries to a list of strings, as confirmed by its consistent usage across the codebase. However, please verify that default rendering (using{{ instructions }}) meets your expectations when multiple strings are provided—even if they are simply concatenated—to ensure the output is formatted as intended. If a more user-friendly display is desired, consider iterating over each instruction or joining them with newlines.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
wren-ai-service/tests/pytest/services/test_ask.py (1)
105-156: Consider adding test coverage for the new retrieval components.While the existing test structure remains functional with the new components, there's no specific test coverage for the
sql_pairs_retrievalandinstructions_retrievalfunctionality. Consider extending the test case to verify these new components are working as expected.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/tests/pytest/services/test_ask.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (1)
wren-ai-service/tests/pytest/services/test_ask.py (1)
55-60: LGTM! Good integration of new retrieval components.The added components (
sql_pairs_retrievalandinstructions_retrieval) align well with the PR objectives to improve SQL pairs and instructions. The implementation correctly initializes both components with parameters from the corresponding pipe_components.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wren-ai-service/tests/pytest/services/mocks.py (2)
16-22: Avoid using mutable default arguments.Using
documents: list = []can lead to unexpected side effects if the list is changed elsewhere. Prefer defaulting toNoneand assigning an empty list internally.Proposed fix:
-class SqlPairsRetrievalMock(retrieval.SqlPairsRetrieval): - def __init__(self, documents: list = []): - self._documents = documents +class SqlPairsRetrievalMock(retrieval.SqlPairsRetrieval): + def __init__(self, documents: Optional[list] = None): + if documents is None: + documents = [] + self._documents = documents
24-29: Same note about mutable default arguments.As with the
SqlPairsRetrievalMock, defaulting to[]can introduce shared state issues. UseNoneinstead.Proposed fix:
-class InstructionsRetrievalMock(retrieval.Instructions): - def __init__(self, documents: list = []): - self._documents = documents +class InstructionsRetrievalMock(retrieval.Instructions): + def __init__(self, documents: Optional[list] = None): + if documents is None: + documents = [] + self._documents = documents
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
wren-ai-service/tests/pytest/services/mocks.py(1 hunks)wren-ai-service/tests/pytest/services/test_ask.py(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/tests/pytest/services/test_ask.py
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wren-ai-service/src/web/v1/services/ask.py (1)
214-214: Add unit test coverage for the new variable.
Defininginstructions = []is fine, but ensure proper usage in the pipeline and test coverage for any logic relying on this variable.wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py (1)
73-76: Consider more specific type annotations
Typing the retrievers asAnyis workable but less precise. Encouraging a more specific interface/type for better maintainability.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
wren-ai-service/src/pipelines/generation/intent_classification.py(7 hunks)wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py(3 hunks)wren-ai-service/src/pipelines/retrieval/retrieval.py(7 hunks)wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py(3 hunks)wren-ai-service/src/web/v1/services/ask.py(7 hunks)wren-ai-service/tests/pytest/services/mocks.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/src/pipelines/generation/intent_classification.py
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (19)
wren-ai-service/src/web/v1/services/ask.py (3)
254-264: Concurrent retrieval is well-structured.
Async gathering forsql_pairs_retrievalandinstructions_retrievalimproves efficiency. No issues found.
266-273: Check for empty output.
Extracting empty lists if nothing is retrieved is handled gracefully. Good defensive coding.
274-334: Intent classification with instructions.
The flow for handlingMISLEADING_QUERYandGENERALlooks correct. Logic is consistent, and concurrency is integrated well.wren-ai-service/tests/pytest/services/mocks.py (4)
12-12: Parameter rename maintains consistency.
Renamingidtoproject_idis clear and consistent. Looks good.
16-22: New mock helps test SQL pairs retrieval.
Implementation is straightforward. This mock is ready for integration tests.
24-30: New mock for instructions retrieval.
Great addition for testing the instructions pipeline.
36-36: Parameter rename in HistoricalQuestionMock.
Aligning with project_id naming. No issues.wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py (3)
38-40: Apply consistent checks.
Usingproject_idas a filter is consistent with changes elsewhere. Ensure tests confirm the filter is correct.
63-64: Consistent embedding usage.
Includingproject_idin the filter logic is appropriate for scoping.
146-146: New run parameter.
Acceptingproject_idaligns with the rest of the pipeline. Straightforward change.wren-ai-service/src/pipelines/retrieval/retrieval.py (4)
123-123: Streamlined list comprehension.
Extractinghistory.questionin one line is neat and clear.
136-136: Parameter rename.
Switchingidtoproject_idis consistent with new naming across modules.
168-168: Consistent filtering for dbschema retrieval.
Adding theproject_idcondition mirrors the pattern used elsewhere.Also applies to: 191-192
479-479: Method signature alignment.
Updatingrunto includeproject_idis consistent with pipeline usage.wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py (5)
44-44: Renaming toproject_idimproves clarity
The parameter rename fromidtoproject_idis more descriptive and aligns well with the filters below.
50-50: Consistent filter usage
Referencing"project_id"here consistently applies the new naming. No issues found.
83-86: Project-based filtering logic looks correct
Using"project_id"in filters and conditionally applying them aligns with the rename. Implementation looks consistent.
171-171: Better param naming
Renaming the parameter toproject_idclarifies intent of the function argument.
177-177: Verify use of empty string fallback
Usingproject_id or ""prevents errors whenproject_idisNone. Confirm that an empty string behaves as intended downstream.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wren-ai-service/tests/pytest/services/mocks.py (2)
16-19: Avoid using a mutable list as a default argument.
Consider usingdocuments: Optional[list] = Noneand initializing with an empty list inside the constructor to prevent shared state across instances.
24-27: Avoid mutable default argument.
Same concern as above; consider changingdocuments: list = []toNoneand check forNoneinside__init__.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/tests/pytest/services/mocks.py(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (6)
wren-ai-service/tests/pytest/services/mocks.py (6)
12-12: Aligning method signature withproject_id.
No issues spotted: this is consistent with the rest of the codebase.
20-21: Mock return structure looks good.
It correctly returns the_documentsin the specified format. No concerns here.
28-29: Returning documents in desired format.
Implementation is straightforward; no additional concerns.
36-36: Project ID parameter updates.
Matches the changes above to unify the parameter name. Implementation is consistent.
47-47: New parameter matches pipeline changes.
The addition ofproject_idin the signature aligns with the broader refactor.
50-51: Optional parameters for SQL and instructions.
These new parameters in the mock help simulate real usage more accurately. Great addition.
2 participants
Summary by CodeRabbit
New Features
Refactor
Chores
These updates deliver a clearer, more informative SQL generation experience with improved responsiveness.