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

wren-ai-service/src/config.py

@@ -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=True)
     max_histories: int = Field(default=5)
     max_sql_correction_retries: int = Field(default=3)
 

wren-ai-service/src/globals.py

@@ -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,
@@ -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(
@@ -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(
@@ -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,
         ),
     )

wren-ai-service/src/pipelines/generation/followup_sql_generation.py

@@ -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
 
@@ -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,
@@ -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,
     )
@@ -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
@@ -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...")
 
@@ -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={
@@ -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,
             },
         )

wren-ai-service/src/pipelines/generation/sql_correction.py

@@ -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.
 
@@ -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:
@@ -46,6 +50,7 @@
 }}
 """
 
+
 sql_correction_user_prompt_template = """
 {% if documents %}
 ### DATABASE SCHEMA ###
@@ -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
@@ -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:

wren-ai-service/src/pipelines/generation/sql_generation.py

@@ -14,13 +14,14 @@
 from src.pipelines.generation.utils.sql import (
     SQL_GENERATION_MODEL_KWARGS,
     SQLGenPostProcessor,
-    calculated_field_instructions,
     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
 
 logger = logging.getLogger("wren-ai-service")
@@ -93,6 +94,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,
@@ -102,10 +104,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,
     )
@@ -157,11 +165,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=sql_generation_user_prompt_template
@@ -189,6 +195,7 @@ async def run(
         use_dry_plan: bool = False,
         allow_dry_plan_fallback: bool = True,
         allow_data_preview: bool = False,
+        sql_knowledge: SqlKnowledge | None = None,
     ):
         logger.info("SQL Generation pipeline is running...")
 
@@ -197,6 +204,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={
@@ -214,6 +226,7 @@ async def run(
                 "allow_dry_plan_fallback": allow_dry_plan_fallback,
                 "data_source": metadata.get("data_source", "local_file"),
                 "allow_data_preview": allow_data_preview,
+                "sql_knowledge": sql_knowledge,
                 **self._components,
             },
         )