feat(wren-ai-service): Add Instructions for SQL Generation

deployment/kustomizations/base/cm.yaml

@@ -177,6 +177,12 @@ data:
       - name: sql_regeneration
         llm: litellm_llm.default
         engine: wren_ui
+      - name: instructions_indexing
+        embedder: litellm_embedder.default
+        document_store: qdrant
+      - name: instructions_retrieval
+        embedder: litellm_embedder.default
+        document_store: qdrant
 
     ---
     settings:
@@ -192,4 +198,9 @@ data:
       langfuse_host: https://cloud.langfuse.com
       langfuse_enable: true
       logging_level: DEBUG
-      development: false
+      development: false
+      historical_question_retrieval_similarity_threshold: 0.9
+      sql_pairs_similarity_threshold: 0.7
+      sql_pairs_retrieval_max_size: 10
+      instructions_similarity_threshold: 0.7
+      instructions_top_k: 10

docker/config.example.yaml

@@ -2,39 +2,39 @@ type: llm
 provider: litellm_llm
 timeout: 120
 models:
-- model: gpt-4o-mini-2024-07-18
-  alias: default
-  api_base: https://api.openai.com/v1
-  kwargs:
-    temperature: 0
-    n: 1
-    # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed
-    seed: 0
-    max_tokens: 4096
-- model: gpt-4o-2024-08-06
-  api_base: https://api.openai.com/v1
-  kwargs:
-    temperature: 0
-    n: 1
-    # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed
-    seed: 0
-    max_tokens: 4096
-- model: o3-mini-2025-01-31
-  api_base: https://api.openai.com/v1
-  kwargs:
-    n: 1
-    seed: 0
-    max_completion_tokens: 4096
-    reasoning_effort: low
+  - model: gpt-4o-mini-2024-07-18
+    alias: default
+    api_base: https://api.openai.com/v1
+    kwargs:
+      temperature: 0
+      n: 1
+      # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed
+      seed: 0
+      max_tokens: 4096
+  - model: gpt-4o-2024-08-06
+    api_base: https://api.openai.com/v1
+    kwargs:
+      temperature: 0
+      n: 1
+      # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed
+      seed: 0
+      max_tokens: 4096
+  - model: o3-mini-2025-01-31
+    api_base: https://api.openai.com/v1
+    kwargs:
+      n: 1
+      seed: 0
+      max_completion_tokens: 4096
+      reasoning_effort: low
 
 ---
 type: embedder
 provider: litellm_embedder
 models:
-- model: text-embedding-3-large
-  alias: default
-  api_base: https://api.openai.com/v1
-  timeout: 120
+  - model: text-embedding-3-large
+    alias: default
+    api_base: https://api.openai.com/v1
+    timeout: 120
 
 ---
 type: engine
@@ -129,6 +129,12 @@ pipes:
   - name: sql_regeneration
     llm: litellm_llm.default
     engine: wren_ui
+  - name: instructions_indexing
+    embedder: litellm_embedder.default
+    document_store: qdrant
+  - name: instructions_retrieval
+    embedder: litellm_embedder.default
+    document_store: qdrant
 
 ---
 settings:
@@ -145,3 +151,8 @@ settings:
   langfuse_enable: true
   logging_level: DEBUG
   development: false
+  historical_question_retrieval_similarity_threshold: 0.9
+  sql_pairs_similarity_threshold: 0.7
+  sql_pairs_retrieval_max_size: 10
+  instructions_similarity_threshold: 0.7
+  instructions_top_k: 10

wren-ai-service/demo/utils.py

@@ -228,33 +228,37 @@ def show_asks_results():
     st.markdown("### Question")
     st.markdown(f"{st.session_state['query']}")
 
-    st.markdown("### Retrieved Tables")
-    retrieved_tables = st.text_input(
-        "Enter the retrieved tables separated by commas, ex: table1, table2, table3",
-        st.session_state["retrieved_tables"],
-        key="retrieved_tables_input",
-    )
+    if st.session_state["asks_results_type"] == "MISLEADING_QUERY":
+        st.markdown(
+            "Misleading query detected. Please try again with a different query."
+        )
+    elif st.session_state["asks_results_type"] == "TEXT_TO_SQL":
+        st.markdown("### Retrieved Tables")
+        retrieved_tables = st.text_input(
+            "Enter the retrieved tables separated by commas, ex: table1, table2, table3",
+            st.session_state["retrieved_tables"],
+            key="retrieved_tables_input",
+        )
 
-    st.markdown("### SQL Generation Reasoning")
-    changed_sql_generation_reasoning = st.text_area(
-        "SQL Generation Reasoning",
-        st.session_state["sql_generation_reasoning"],
-        key="sql_generation_reasoning_input",
-        height=250,
-        on_change=on_change_sql_generation_reasoning,
-    )
+        st.markdown("### SQL Generation Reasoning")
+        changed_sql_generation_reasoning = st.text_area(
+            "SQL Generation Reasoning",
+            st.session_state["sql_generation_reasoning"],
+            key="sql_generation_reasoning_input",
+            height=250,
+            on_change=on_change_sql_generation_reasoning,
+        )
 
-    st.button(
-        "Regenerate SQL",
-        on_click=on_click_regenerate_sql,
-        args=(
-            retrieved_tables.split(", "),
-            changed_sql_generation_reasoning,
-        ),
-    )
+        st.button(
+            "Regenerate SQL",
+            on_click=on_click_regenerate_sql,
+            args=(
+                retrieved_tables.split(", "),
+                changed_sql_generation_reasoning,
+            ),
+        )
 
-    st.markdown("### SQL Query Result")
-    if st.session_state["asks_results_type"] == "TEXT_TO_SQL":
+        st.markdown("### SQL Query Result")
         edited_sql = st.text_area(
             label="SQL Query Result",
             value=sqlparse.format(
@@ -283,10 +287,6 @@ def show_asks_results():
         st.session_state["asks_details_result"] = None
         st.session_state["preview_data_button_index"] = None
         st.session_state["preview_sql"] = None
-    elif st.session_state["asks_results_type"] == "MISLEADING_QUERY":
-        st.markdown(
-            "Misleading query detected. Please try again with a different query."
-        )
 
 
 def show_asks_details_results():

wren-ai-service/docs/config_examples/config.azure.yaml

@@ -132,6 +132,12 @@ pipes:
   - name: sql_regeneration
     llm: litellm_llm.default
     engine: wren_ui
+  - name: instructions_indexing
+    embedder: litellm_embedder.default
+    document_store: qdrant
+  - name: instructions_retrieval
+    embedder: litellm_embedder.default
+    document_store: qdrant
 
 ---
 settings:
@@ -147,3 +153,8 @@ settings:
   langfuse_enable: true
   logging_level: DEBUG
   development: false
+  historical_question_retrieval_similarity_threshold: 0.9
+  sql_pairs_similarity_threshold: 0.7
+  sql_pairs_retrieval_max_size: 10
+  instructions_similarity_threshold: 0.7
+  instructions_top_k: 10

wren-ai-service/docs/config_examples/config.deepseek.yaml

@@ -142,6 +142,12 @@ pipes:
   - name: sql_regeneration
     llm: litellm_llm.default
     engine: wren_ui
+  - name: instructions_indexing
+    embedder: litellm_embedder.default
+    document_store: qdrant
+  - name: instructions_retrieval
+    embedder: litellm_embedder.default
+    document_store: qdrant
 
 ---
 settings:
@@ -157,3 +163,8 @@ settings:
   langfuse_enable: true
   logging_level: DEBUG
   development: true
+  historical_question_retrieval_similarity_threshold: 0.9
+  sql_pairs_similarity_threshold: 0.7
+  sql_pairs_retrieval_max_size: 10
+  instructions_similarity_threshold: 0.7
+  instructions_top_k: 10

wren-ai-service/docs/config_examples/config.google_ai_studio.yaml

@@ -131,6 +131,12 @@ pipes:
   - name: sql_regeneration
     llm: litellm_llm.default
     engine: wren_ui
+  - name: instructions_indexing
+    embedder: litellm_embedder.default
+    document_store: qdrant
+  - name: instructions_retrieval
+    embedder: litellm_embedder.default
+    document_store: qdrant
 
 ---
 settings:
@@ -146,3 +152,8 @@ settings:
   langfuse_enable: true
   logging_level: DEBUG
   development: true
+  historical_question_retrieval_similarity_threshold: 0.9
+  sql_pairs_similarity_threshold: 0.7
+  sql_pairs_retrieval_max_size: 10
+  instructions_similarity_threshold: 0.7
+  instructions_top_k: 10

wren-ai-service/docs/config_examples/config.groq.yaml

@@ -124,6 +124,12 @@ pipes:
   - name: sql_regeneration
     llm: litellm_llm.default
     engine: wren_ui
+  - name: instructions_indexing
+    embedder: litellm_embedder.default
+    document_store: qdrant
+  - name: instructions_retrieval
+    embedder: litellm_embedder.default
+    document_store: qdrant
 
 ---
 settings:
@@ -139,3 +145,8 @@ settings:
   langfuse_enable: true
   logging_level: DEBUG
   development: true
+  historical_question_retrieval_similarity_threshold: 0.9
+  sql_pairs_similarity_threshold: 0.7
+  sql_pairs_retrieval_max_size: 10
+  instructions_similarity_threshold: 0.7
+  instructions_top_k: 10

wren-ai-service/docs/config_examples/config.ollama.yaml

@@ -121,6 +121,12 @@ pipes:
   - name: sql_regeneration
     llm: litellm_llm.default
     engine: wren_ui
+  - name: instructions_indexing
+    embedder: litellm_embedder.default
+    document_store: qdrant
+  - name: instructions_retrieval
+    embedder: litellm_embedder.default
+    document_store: qdrant
 
 ---
 settings:
@@ -136,3 +142,8 @@ settings:
   langfuse_enable: true
   logging_level: DEBUG
   development: true
+  historical_question_retrieval_similarity_threshold: 0.9
+  sql_pairs_similarity_threshold: 0.7
+  sql_pairs_retrieval_max_size: 10
+  instructions_similarity_threshold: 0.7
+  instructions_top_k: 10

wren-ai-service/eval/__init__.py

@@ -11,6 +11,7 @@ class EvalSettings(Settings):
     config_path: str = "eval/config.yaml"
     openai_api_key: SecretStr = Field(alias="OPENAI_API_KEY")
     allow_sql_samples: bool = True
+    allow_instructions: bool = True
     db_path_for_duckdb: str = ""
 
     # BigQuery

wren-ai-service/eval/evaluation.py

@@ -171,7 +171,7 @@ def _average_score(self, meta: dict) -> None:
 
     dataset = parse_toml(meta["evaluation_dataset"])
     metrics = pipelines.metrics_initiator(
-        meta["pipeline"], dataset, pipe_components, args.semantics
+        meta["pipeline"], dataset, pipe_components, args.semantics, settings
     )
 
     evaluator = Evaluator(**metrics)