fix(wren-ai-service): fix sql expansion latency and retrieval issue

[cyyeh]

Summary by CodeRabbit

• New Features
  • Enhanced service configuration with a flexible approach to handling data sources.
  • Improved retrieval operations by incorporating context from previous interactions for more relevant outcomes.

[coderabbitai]

Walkthrough

This pull request integrates additional parameters into the retrieval configuration across two modules. In the service container initialization, a new flag from the settings controls whether database schemas may be used without pruning. In the SQL expansion service, a historical context parameter is added to the retrieval run call. These changes adjust how retrieval components are instantiated and how historical data is passed without altering the overall function flow.

Changes

File	Change Summary
wren-ai-service/.../globals.py	Added the allow_using_db_schemas_without_pruning parameter (sourced from settings) to the instantiation of the retrieval.Retrieval class in both ask_service and sql_expansion_service.
wren-ai-service/.../sql_expansion.py	Introduced the histories parameter in the retrieval pipeline’s run method, passing a list containing the extracted history from the request.

Sequence Diagram(s)

sequenceDiagram
    participant Container as Service Container
    participant Settings
    participant Retrieval
    Container->>Settings: Retrieve allow_using_db_schemas_without_pruning
    Container->>Retrieval: Instantiate ask_service & sql_expansion_service with new flag

Loading

sequenceDiagram
    participant Request
    participant SQLService
    participant Retrieval

    Request->>SQLService: Send SQL expansion request with histories
    SQLService->>Retrieval: Call run(histories=[history])

Loading

Possibly related PRs

• chore(wren-ai-service): minor updates #1389: Enhances retrieval configuration by exposing the historical_question_retrieval_similarity_threshold parameter, aligning with the new retrieval parameter changes.
• chore(wren-ai-service): add followup sql generation reasoning #1407: Introduces a new service for follow-up SQL generation reasoning, which relates to the modifications in the retrieval setup of this PR.

Suggested labels

wren-ai-service

Suggested reviewers

• paopa

Poem

I'm a bunny in the code burrow so deep,
Hopping through changes while others sleep.
Schemas and histories now neatly aligned,
In our code garden, new features we find.
Leaping forward with each joyful leap!

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 7f2d940 and 0a4a64b.

📒 Files selected for processing (2)

• wren-ai-service/src/globals.py (1 hunks)
• wren-ai-service/src/web/v1/services/sql_expansion.py (1 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 (2)

wren-ai-service/src/web/v1/services/sql_expansion.py (1)

135-137:

Details

❓ Verification inconclusive

Good enhancement: Passing historical context to the retrieval pipeline.

Adding the histories=[history] parameter to the retrieval run will help provide historical context to the retrieval process, potentially improving the relevance of retrieved documents.

Let's verify if the retrieval pipeline properly handles this new parameter:

---

🏁 Script executed:

#!/bin/bash
# Check if the retrieval pipeline handles the 'histories' parameter
rg -A 3 -B 3 "def run\(" src/pipelines/retrieval/

Length of output: 116

---

Manual Verification Required: Check Retrieval Pipeline Parameter Handling

The new parameter (histories=[history]) seems correctly added to the retrieval call in wren-ai-service/src/web/v1/services/sql_expansion.py. However, the automated search for a corresponding run function within a src/pipelines/retrieval/ directory did not locate any matching files. This might indicate that the retrieval pipeline is implemented elsewhere or has been relocated. Please manually verify that the retrieval pipeline (wherever it is defined) appropriately handles the histories parameter.

• Action Required: Confirm manually that the added historical context is processed correctly by the retrieval pipeline.

wren-ai-service/src/globals.py (1)

198-198:

Details

❓ Verification inconclusive

Consistency improvement: Added missing configuration parameter.

Adding the allow_using_db_schemas_without_pruning parameter to the SQL expansion service's retrieval configuration ensures consistency with other services and should help reduce latency by potentially skipping an unnecessary pruning step.

Let's verify that this parameter is properly defined in the settings:

---

🏁 Script executed:

#!/bin/bash
# Check if the parameter is defined in settings
rg "allow_using_db_schemas_without_pruning" src/config/

Length of output: 107

---

#!/bin/bash
# Search for the parameter definition across the entire repository
rg "allow_using_db_schemas_without_pruning" -n .

The initial search in the expected "src/config/" directory failed because that directory doesn't exist. To ensure the parameter is properly defined (likely in a different settings file such as wren-ai-service/src/settings.py), please run the above script and verify the parameter's definition across the entire repository.

Action Required: Verify the definition and location of the configuration parameter.

• Confirm that allow_using_db_schemas_without_pruning is defined in the settings module (e.g., in wren-ai-service/src/settings.py or another relevant configuration file).
• If the definition is missing or located in an unexpected file, update either the settings or the globals reference accordingly.

✨ Finishing Touches

• 📝 Generate Docstrings

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

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)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai plan to trigger planning for file edits and PR creation.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.