chore(wren-ai-service): minor updates #1293
Conversation
cyyeh
added
module/ai-service
WalkthroughThis pull request introduces multiple modifications. A new dependency ( Changes
Sequence Diagram(s)sequenceDiagram
participant User as User
participant AS as AskService
participant RM as Retrieval Module
participant SQ as SQLGeneration Pipeline
User->>AS: Send ask request
AS->>RM: Invoke retrieval (e.g., check_using_db_schemas_without_pruning)
RM-->>AS: Return table details (name & DDL)
AS->>AS: Extract table names and aggregate table DDLs
AS->>SQ: Forward request with table_ddls context
SQ-->>AS: Return SQL response
AS->>User: Respond with AskResultResponse (includes retrieved_tables)
Suggested reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms (4)
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: 2
🔭 Outside diff range comments (3)
wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py (3)
41-47:⚠️ Potential issueFix test assertion to include 'columns' field.
The test is failing because the document content assertion is missing the 'columns' field that was added in the implementation.
Update the assertion to include the 'columns' field:
assert document.content == str( { "name": "user", - "mdl_type": "MODEL", "description": "A table containing user information.", + "columns": "" } )🧰 Tools
🪛 GitHub Actions: AI Service Test
[error] 41-41: Assertion error: The actual content of the document does not match the expected content. The 'columns' field is missing from the actual content.
76-82:⚠️ Potential issueFix test assertion to include 'columns' field.
Similar to the previous test, this assertion is missing the 'columns' field.
Update the assertion to include the 'columns' field:
assert document_1.content == str( { "name": "user", - "mdl_type": "MODEL", "description": "A table containing user information.", + "columns": "" } )🧰 Tools
🪛 GitHub Actions: AI Service Test
[error] 76-76: Assertion error: The actual content of the document does not match the expected content. The 'columns' field is missing from the actual content.
86-92:⚠️ Potential issueFix test assertion to include 'columns' field.
The assertion for document_2 also needs to be updated.
Update the assertion to include the 'columns' field:
assert document_2.content == str( { "name": "order", - "mdl_type": "MODEL", "description": "A table containing order details.", + "columns": "" } )
🧹 Nitpick comments (1)
wren-ai-service/src/pipelines/generation/sql_answer.py (1)
34-34: Fix typo in output format instruction.There's a missing space in "stringformat".
-Please provide your response in proper Markdown stringformat. +Please provide your response in proper Markdown string format.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
wren-ai-service/poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (7)
wren-ai-service/pyproject.toml(1 hunks)wren-ai-service/src/pipelines/generation/sql_answer.py(1 hunks)wren-ai-service/src/pipelines/indexing/table_description.py(3 hunks)wren-ai-service/src/pipelines/retrieval/retrieval.py(2 hunks)wren-ai-service/src/web/v1/services/ask.py(11 hunks)wren-ai-service/src/web/v1/services/question_recommendation.py(2 hunks)wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py(4 hunks)
🧰 Additional context used
🪛 GitHub Actions: AI Service Test
wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py
[error] 41-41: Assertion error: The actual content of the document does not match the expected content. The 'columns' field is missing from the actual content.
[error] 76-76: Assertion error: The actual content of the document does not match the expected content. The 'columns' field is missing from the actual content.
[error] 126-126: Assertion error: The actual content of the document does not match the expected content. The 'columns' field is missing from the actual content.
⏰ 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/pipelines/generation/sql_answer.py (1)
30-30: LGTM!The instruction to prevent markdown code block syntax is clear and appropriate.
wren-ai-service/src/pipelines/indexing/table_description.py (2)
34-34: LGTM!Adding the table name to document metadata improves searchability and context.
57-57: LGTM!The changes to include column information and format it as a comma-separated string enhance the table description's completeness.
Also applies to: 71-71
wren-ai-service/src/web/v1/services/question_recommendation.py (1)
79-79: LGTM!The changes to use table_ddls instead of documents align with the PR objectives and improve the clarity of what data is being passed to the SQL generation pipeline.
Also applies to: 87-87, 97-97
wren-ai-service/src/pipelines/retrieval/retrieval.py (4)
229-234: LGTM! Improved data structure to include table names.The change enhances the data structure by including both table name and DDL, making it more informative and easier to track the source of each DDL.
241-246: LGTM! Consistent data structure across table types.The same structured format is applied to metrics and views, maintaining consistency in how table information is returned.
Also applies to: 249-254
256-259: LGTM! Updated token count calculation.The token count calculation is correctly updated to work with the new dictionary structure.
349-354: LGTM! Consistent data structure in filtered results.The structured format is maintained when filtering results based on columns and tables needed.
Also applies to: 361-366, 369-374
wren-ai-service/src/web/v1/services/ask.py (4)
97-97: LGTM! Added retrieved_tables field to response model.The new optional field allows tracking which tables were used in generating the response.
314-315: LGTM! Extracted table information from retrieval results.The code correctly extracts both table names and DDLs from the new structured retrieval results.
344-344: LGTM! Consistent propagation of table names.Table names are consistently included in all response states (planning, generating, correcting, finished, failed).
Also applies to: 372-372, 382-382, 440-440, 478-478, 495-495
358-358: LGTM! Updated context passing.The code correctly passes table DDLs as contexts to various pipeline stages.
Also applies to: 396-396, 410-410, 446-446
Summary by CodeRabbit
New Features
Bug Fixes