chore(wren-ai-service): minor updates

deployment/kustomizations/base/cm.yaml

@@ -184,6 +184,8 @@ data:
       column_indexing_batch_size: 50
       table_retrieval_size: 10
       table_column_retrieval_size: 100
+      allow_intent_classification: true
+      allow_sql_generation_reasoning: true
       allow_using_db_schemas_without_pruning: false
       query_cache_maxsize: 1000
       query_cache_ttl: 3600

docker/config.example.yaml

@@ -136,6 +136,8 @@ settings:
   column_indexing_batch_size: 50
   table_retrieval_size: 10
   table_column_retrieval_size: 100
+  allow_intent_classification: true
+  allow_sql_generation_reasoning: true
   allow_using_db_schemas_without_pruning: false
   query_cache_maxsize: 1000
   query_cache_ttl: 3600

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

@@ -139,6 +139,8 @@ settings:
   table_retrieval_size: 10
   table_column_retrieval_size: 100
   allow_using_db_schemas_without_pruning: false
+  allow_intent_classification: true
+  allow_sql_generation_reasoning: true  
   query_cache_maxsize: 1000
   query_cache_ttl: 3600
   langfuse_host: https://cloud.langfuse.com

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

@@ -149,6 +149,8 @@ settings:
   table_retrieval_size: 10
   table_column_retrieval_size: 100
   allow_using_db_schemas_without_pruning: false # if you want to use db schemas without pruning, set this to true. It will be faster
+  allow_intent_classification: true
+  allow_sql_generation_reasoning: true  
   query_cache_maxsize: 1000
   query_cache_ttl: 3600
   langfuse_host: https://cloud.langfuse.com

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

@@ -138,6 +138,8 @@ settings:
   table_retrieval_size: 10
   table_column_retrieval_size: 100
   allow_using_db_schemas_without_pruning: false # if you want to use db schemas without pruning, set this to true. It will be faster
+  allow_intent_classification: true
+  allow_sql_generation_reasoning: true  
   query_cache_maxsize: 1000
   query_cache_ttl: 3600
   langfuse_host: https://cloud.langfuse.com

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

@@ -131,6 +131,8 @@ settings:
   table_retrieval_size: 10
   table_column_retrieval_size: 100
   allow_using_db_schemas_without_pruning: false # if you want to use db schemas without pruning, set this to true. It will be faster
+  allow_intent_classification: true
+  allow_sql_generation_reasoning: true  
   query_cache_maxsize: 1000
   query_cache_ttl: 3600
   langfuse_host: https://cloud.langfuse.com

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

@@ -128,6 +128,8 @@ settings:
   table_retrieval_size: 10
   table_column_retrieval_size: 100
   allow_using_db_schemas_without_pruning: false # if you want to use db schemas without pruning, set this to true. It will be faster
+  allow_intent_classification: true
+  allow_sql_generation_reasoning: true
   query_cache_maxsize: 1000
   query_cache_ttl: 3600
   langfuse_host: https://cloud.langfuse.com

wren-ai-service/src/config.py

@@ -30,6 +30,7 @@ class Settings(BaseSettings):
     table_retrieval_size: int = Field(default=10)
     table_column_retrieval_size: int = Field(default=100)
     allow_using_db_schemas_without_pruning: bool = Field(default=False)
+    historical_question_retrieval_similarity_threshold: float = Field(default=0.9)
     sql_pairs_similarity_threshold: float = Field(default=0.7)
     sql_pairs_retrieval_max_size: int = Field(default=10)
 

wren-ai-service/src/globals.py

@@ -86,6 +86,7 @@ def create_service_container(
                 ),
                 "historical_question": retrieval.HistoricalQuestionRetrieval(
                     **pipe_components["historical_question_retrieval"],
+                    historical_question_retrieval_similarity_threshold=settings.historical_question_retrieval_similarity_threshold,
                 ),
                 "sql_pairs_retrieval": retrieval.SqlPairsRetrieval(
                     **pipe_components["sql_pairs_retrieval"],

wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py

@@ -102,9 +102,16 @@ async def retrieval(
 
 
 @observe(capture_input=False)
-def filtered_documents(retrieval: dict, score_filter: ScoreFilter) -> dict:
+def filtered_documents(
+    retrieval: dict,
+    score_filter: ScoreFilter,
+    historical_question_retrieval_similarity_threshold: float,
+) -> dict:
     if retrieval:
-        return score_filter.run(documents=retrieval.get("documents"), score=0.9)
+        return score_filter.run(
+            documents=retrieval.get("documents"),
+            score=historical_question_retrieval_similarity_threshold,
+        )
 
     return {}
 
@@ -127,6 +134,7 @@ def __init__(
         self,
         embedder_provider: EmbedderProvider,
         document_store_provider: DocumentStoreProvider,
+        historical_question_retrieval_similarity_threshold: Optional[float] = 0.9,
         **kwargs,
     ) -> None:
         view_questions_store = document_store_provider.get_store(
@@ -148,6 +156,10 @@ def __init__(
             "output_formatter": OutputFormatter(),
         }
 
+        self._configs = {
+            "historical_question_retrieval_similarity_threshold": historical_question_retrieval_similarity_threshold,
+        }
+
         super().__init__(
             AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
         )
@@ -161,6 +173,7 @@ async def run(self, query: str, id: Optional[str] = None):
                 "query": query,
                 "id": id or "",
                 **self._components,
+                **self._configs,
             },
         )
 

wren-ai-service/tools/config/config.example.yaml

@@ -154,6 +154,8 @@ settings:
   column_indexing_batch_size: 50
   table_retrieval_size: 10
   table_column_retrieval_size: 100
+  allow_intent_classification: true
+  allow_sql_generation_reasoning: true
   allow_using_db_schemas_without_pruning: false
   query_cache_maxsize: 1000
   query_cache_ttl: 3600

wren-ai-service/tools/config/config.full.yaml

@@ -155,6 +155,8 @@ settings:
   table_retrieval_size: 10
   table_column_retrieval_size: 100
   query_cache_maxsize: 1000
+  allow_intent_classification: true
+  allow_sql_generation_reasoning: true
   allow_using_db_schemas_without_pruning: false
   query_cache_ttl: 3600
   langfuse_host: https://cloud.langfuse.com