Skip to content

feat(wren-ai-service): get knowledge from ibis (ai-env-changed) #2043

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cyyeh merged 19 commits into from
Dec 10, 2025
Merged

feat(wren-ai-service): get knowledge from ibis (ai-env-changed) #2043

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
9f9a5da
Refactor SQL generation prompts and rules in the AI service. Moved th…
yichieh-lu Nov 17, 2025
3917330
Update environment versions for WREN engine, IBIS server, and WREN UI…
yichieh-lu Nov 17, 2025
23b25c7
Add async method to retrieve SQL knowledge from WREN API
yichieh-lu Nov 17, 2025
c09d4b3
Refactor SQL generation system prompt for improved clarity and consis…
yichieh-lu Nov 17, 2025
22accf1
Update SQL generation reasoning system prompt for enhanced clarity an…
yichieh-lu Nov 17, 2025
96e1fff
Add SQL knowledge retrieval to service container
yichieh-lu Nov 17, 2025
afe2166
Refactor SQL generation components to utilize getter functions for in…
yichieh-lu Nov 17, 2025
7abf4b0
Add SQL knowledge retrieval configuration to settings and service con…
yichieh-lu Nov 17, 2025
b5029fb
Add support for SQL knowledge retrieval in multiple services
yichieh-lu Nov 17, 2025
5f548db
Implement SQL knowledge retrieval pipeline and refactor SQL generatio…
yichieh-lu Nov 17, 2025
25b6bd0
Remove generator initialization from SQL generation classes to stream…
yichieh-lu Nov 17, 2025
a129c8f
Enhance SQL generation and correction pipelines with SqlKnowledge int…
yichieh-lu Nov 17, 2025
fceeb3e
Refactor SQLGeneration class to remove print statement and enhance co…
yichieh-lu Nov 17, 2025
aa1b5cc
update
cyyeh Nov 18, 2025
844a6b5
update
cyyeh Nov 18, 2025
af3ec91
fix
cyyeh Nov 19, 2025
81c1967
update
cyyeh Nov 19, 2025
49b93a7
Merge branch 'main' into feat/wren-ai-service/get-knowledge-from-ibis
cyyeh Dec 10, 2025
4fa6c11
update configs
cyyeh Dec 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions docker/config.example.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
Comment on lines +166 to +168
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Add missing allow_sql_knowledge_retrieval configuration flag.

Per the PR objectives, a new allow_sql_knowledge_retrieval flag (default: true) controls whether the pipeline retrieves SQL knowledge or falls back to Ibis prompts. This flag is not present in the settings section, but it should be added to align with the feature's design.

Add the configuration flag to the settings section:

 settings:
   doc_endpoint: https://docs.getwren.ai
   is_oss: true
   engine_timeout: 30
   column_indexing_batch_size: 50
   table_retrieval_size: 10
   table_column_retrieval_size: 100
   allow_intent_classification: true
   allow_sql_generation_reasoning: true
   allow_sql_functions_retrieval: true
+  allow_sql_knowledge_retrieval: true
   enable_column_pruning: false
   max_sql_correction_retries: 3

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents 
In docker/config.example.yaml around lines 166 to 168, the settings block is
missing the new allow_sql_knowledge_retrieval boolean flag; add
allow_sql_knowledge_retrieval: true under the settings section (aligning with
defaults) so the pipeline can toggle SQL knowledge retrieval vs Ibis prompts,
and ensure it is documented/commented the same way as nearby flags for clarity.


---
settings:
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.anthropic.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.azure.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.bedrock.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.deepseek.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.google_vertexai.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.grok.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.groq.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.lm_studio.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.ollama.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.open_router.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.qwen3.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
4 changes: 3 additions & 1 deletion wren-ai-service/docs/config_examples/config.zhipu.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,9 @@ pipes:
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default

- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
Expand Down
1 change: 1 addition & 0 deletions wren-ai-service/src/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ class Settings(BaseSettings):
allow_sql_generation_reasoning: bool = Field(default=True)
allow_sql_functions_retrieval: bool = Field(default=True)
allow_sql_diagnosis: bool = Field(default=True)
allow_sql_knowledge_retrieval: bool = Field(default=False)
max_histories: int = Field(default=5)
max_sql_correction_retries: int = Field(default=3)

Expand Down
5 changes: 4 additions & 1 deletion wren-ai-service/src/force_update_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,10 @@ def update_config():
# Update engine name in all pipelines
for pipe in doc.get("pipes", []):
if "engine" in pipe:
if pipe["name"] == "sql_functions_retrieval":
if pipe["name"] in [
"sql_functions_retrieval",
"sql_knowledge_retrieval",
]:
pipe["engine"] = "wren_ibis"
else:
pipe["engine"] = "wren_ui"
Expand Down
16 changes: 16 additions & 0 deletions wren-ai-service/src/globals.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,11 +150,15 @@ def create_service_container(
),
"sql_functions_retrieval": _sql_functions_retrieval_pipeline,
"sql_diagnosis": _sql_diagnosis_pipeline,
"sql_knowledge_retrieval": retrieval.SqlKnowledges(
**pipe_components["sql_knowledge_retrieval"],
),
},
allow_intent_classification=settings.allow_intent_classification,
allow_sql_generation_reasoning=settings.allow_sql_generation_reasoning,
allow_sql_functions_retrieval=settings.allow_sql_functions_retrieval,
allow_sql_diagnosis=settings.allow_sql_diagnosis,
allow_sql_knowledge_retrieval=settings.allow_sql_knowledge_retrieval,
max_histories=settings.max_histories,
enable_column_pruning=settings.enable_column_pruning,
max_sql_correction_retries=settings.max_sql_correction_retries,
Expand All @@ -171,9 +175,13 @@ def create_service_container(
),
"sql_correction": _sql_correction_pipeline,
"sql_diagnosis": _sql_diagnosis_pipeline,
"sql_knowledge_retrieval": retrieval.SqlKnowledges(
**pipe_components["sql_knowledge_retrieval"],
),
},
allow_sql_functions_retrieval=settings.allow_sql_functions_retrieval,
allow_sql_diagnosis=settings.allow_sql_diagnosis,
allow_sql_knowledge_retrieval=settings.allow_sql_knowledge_retrieval,
**query_cache,
),
chart_service=services.ChartService(
Expand Down Expand Up @@ -225,8 +233,12 @@ def create_service_container(
"sql_pairs_retrieval": _sql_pair_retrieval_pipeline,
"instructions_retrieval": _instructions_retrieval_pipeline,
"sql_functions_retrieval": _sql_functions_retrieval_pipeline,
"sql_knowledge_retrieval": retrieval.SqlKnowledges(
**pipe_components["sql_knowledge_retrieval"],
),
},
allow_sql_functions_retrieval=settings.allow_sql_functions_retrieval,
allow_sql_knowledge_retrieval=settings.allow_sql_knowledge_retrieval,
**query_cache,
),
sql_pairs_service=services.SqlPairsService(
Expand Down Expand Up @@ -256,7 +268,11 @@ def create_service_container(
),
"db_schema_retrieval": _db_schema_retrieval_pipeline,
"sql_correction": _sql_correction_pipeline,
"sql_knowledge_retrieval": retrieval.SqlKnowledges(
**pipe_components["sql_knowledge_retrieval"],
),
},
allow_sql_knowledge_retrieval=settings.allow_sql_knowledge_retrieval,
**query_cache,
),
)
Expand Down
35 changes: 24 additions & 11 deletions wren-ai-service/src/pipelines/generation/followup_sql_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,14 +14,15 @@
from src.pipelines.generation.utils.sql import (
SQL_GENERATION_MODEL_KWARGS,
SQLGenPostProcessor,
calculated_field_instructions,
construct_ask_history_messages,
construct_instructions,
json_field_instructions,
metric_instructions,
sql_generation_system_prompt,
get_calculated_field_instructions,
get_json_field_instructions,
get_metric_instructions,
get_sql_generation_system_prompt,
)
from src.pipelines.retrieval.sql_functions import SqlFunction
from src.pipelines.retrieval.sql_knowledge import SqlKnowledge
from src.utils import trace_cost
from src.web.v1.services.ask import AskHistory

Expand Down Expand Up @@ -97,6 +98,7 @@ def prompt(
has_metric: bool = False,
has_json_field: bool = False,
sql_functions: list[SqlFunction] | None = None,
sql_knowledge: SqlKnowledge | None = None,
) -> dict:
_prompt = prompt_builder.run(
query=query,
Expand All @@ -106,10 +108,16 @@ def prompt(
instructions=instructions,
),
calculated_field_instructions=(
calculated_field_instructions if has_calculated_field else ""
get_calculated_field_instructions(sql_knowledge)
if has_calculated_field
else ""
),
metric_instructions=(
get_metric_instructions(sql_knowledge) if has_metric else ""
),
json_field_instructions=(
get_json_field_instructions(sql_knowledge) if has_json_field else ""
),
metric_instructions=(metric_instructions if has_metric else ""),
json_field_instructions=(json_field_instructions if has_json_field else ""),
sql_samples=sql_samples,
sql_functions=sql_functions,
)
Expand Down Expand Up @@ -160,11 +168,9 @@ def __init__(
document_store_provider.get_store("project_meta")
)

self._llm_provider = llm_provider

self._components = {
"generator": llm_provider.get_generator(
system_prompt=sql_generation_system_prompt,
generation_kwargs=SQL_GENERATION_MODEL_KWARGS,
),
"generator_name": llm_provider.get_model(),
"prompt_builder": PromptBuilder(
template=text_to_sql_with_followup_user_prompt_template
Expand Down Expand Up @@ -192,6 +198,7 @@ async def run(
sql_functions: list[SqlFunction] | None = None,
use_dry_plan: bool = False,
allow_dry_plan_fallback: bool = True,
sql_knowledge: SqlKnowledge | None = None,
):
logger.info("Follow-Up SQL Generation pipeline is running...")

Expand All @@ -200,6 +207,11 @@ async def run(
else:
metadata = {}

self._components["generator"] = self._llm_provider.get_generator(
system_prompt=get_sql_generation_system_prompt(sql_knowledge),
generation_kwargs=SQL_GENERATION_MODEL_KWARGS,
)

return await self._pipe.execute(
["post_process"],
inputs={
Expand All @@ -217,6 +229,7 @@ async def run(
"use_dry_plan": use_dry_plan,
"allow_dry_plan_fallback": allow_dry_plan_fallback,
"data_source": metadata.get("data_source", "local_file"),
"sql_knowledge": sql_knowledge,
**self._components,
},
)
22 changes: 15 additions & 7 deletions wren-ai-service/src/pipelines/generation/sql_correction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,21 @@
from src.pipelines.common import clean_up_new_lines, retrieve_metadata
from src.pipelines.generation.utils.sql import (
SQL_GENERATION_MODEL_KWARGS,
TEXT_TO_SQL_RULES,
SQLGenPostProcessor,
construct_instructions,
get_text_to_sql_rules,
)
from src.pipelines.retrieval.sql_functions import SqlFunction
from src.pipelines.retrieval.sql_knowledge import SqlKnowledge
from src.utils import trace_cost

logger = logging.getLogger("wren-ai-service")


sql_correction_system_prompt = f"""
def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) -> str:
text_to_sql_rules = get_text_to_sql_rules(sql_knowledge)

return f"""
### TASK ###
You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect ANSI SQL query.

Expand All @@ -36,7 +40,7 @@
### SQL RULES ###
Make sure you follow the SQL Rules strictly.

{TEXT_TO_SQL_RULES}
{text_to_sql_rules}

### FINAL ANSWER FORMAT ###
The final answer must be in JSON format:
Expand All @@ -46,6 +50,7 @@
}}
"""


sql_correction_user_prompt_template = """
{% if documents %}
### DATABASE SCHEMA ###
Expand Down Expand Up @@ -136,12 +141,9 @@ def __init__(
self._retriever = document_store_provider.get_retriever(
document_store_provider.get_store("project_meta")
)
self._llm_provider = llm_provider

self._components = {
"generator": llm_provider.get_generator(
system_prompt=sql_correction_system_prompt,
generation_kwargs=SQL_GENERATION_MODEL_KWARGS,
),
"generator_name": llm_provider.get_model(),
"prompt_builder": PromptBuilder(
template=sql_correction_user_prompt_template
Expand All @@ -163,9 +165,15 @@ async def run(
project_id: str | None = None,
use_dry_plan: bool = False,
allow_dry_plan_fallback: bool = True,
sql_knowledge: SqlKnowledge | None = None,
):
logger.info("SQLCorrection pipeline is running...")

self._components["generator"] = self._llm_provider.get_generator(
system_prompt=get_sql_correction_system_prompt(sql_knowledge),
generation_kwargs=SQL_GENERATION_MODEL_KWARGS,
)

if use_dry_plan:
metadata = await retrieve_metadata(project_id or "", self._retriever)
else:
Expand Down
Loading
Loading