feat(wren-ai-service): sql pairs

deployment/kustomizations/base/cm.yaml

@@ -152,6 +152,16 @@ data:
         document_store: qdrant
       - name: data_assistance
         llm: litellm_llm.gpt-4o-mini
+      - name: sql_pairs_preparation
+        document_store: qdrant
+        embedder: openai_embedder.text-embedding-3-large
+        llm: litellm_llm.gpt-4o-mini
+      - name: sql_pairs_deletion
+        document_store: qdrant
+        embedder: openai_embedder.text-embedding-3-large
+      - name: sql_pairs_retrieval
+        document_store: qdrant
+        embedder: openai_embedder.text-embedding-3-large
       - name: sql_executor
         engine: wren_ui
       - name: chart_generation
@@ -168,6 +178,5 @@ data:
       query_cache_ttl: 3600
       langfuse_host: https://cloud.langfuse.com
       langfuse_enable: true
-      enable_timer: false
       logging_level: DEBUG
       development: false

docker/config.example.yaml

@@ -100,6 +100,17 @@ pipes:
     document_store: qdrant
   - name: data_assistance
     llm: litellm_llm.gpt-4o-mini
+  - document_store: qdrant
+    embedder: openai_embedder.text-embedding-3-large
+    llm: litellm_llm.gpt-4o-mini
+    name: sql_pairs_preparation
+  - document_store: qdrant
+    embedder: openai_embedder.text-embedding-3-large
+    name: sql_pairs_deletion
+  - document_store: qdrant
+    embedder: openai_embedder.text-embedding-3-large
+    name: sql_pairs_retrieval
+    llm: litellm_llm.gpt-4o-mini
   - name: preprocess_sql_data
     llm: litellm_llm.gpt-4o-mini
   - name: sql_executor
@@ -118,6 +129,5 @@ settings:
   query_cache_ttl: 3600
   langfuse_host: https://cloud.langfuse.com
   langfuse_enable: true
-  enable_timer: false
   logging_level: DEBUG
   development: false

wren-ai-service/README.md

@@ -98,10 +98,6 @@ For a comprehensive understanding of how to evaluate the pipelines, please refer
 
 ### Estimate the Speed of the Pipeline
 
-- to evaluate the speed of the pipeline, you can enable the timer
-  - add environment variables `ENABLE_TIMER=1` in `.env.dev`
-  - restart wren ai service
-  - check the logs in the terminal
 - to run the load test
   - setup `DATASET_NAME` in `.env.dev`
   - adjust test config if needed

wren-ai-service/docs/configuration.md

@@ -134,7 +134,6 @@ The configuration file (`config.yaml`) is structured into several sections, each
      query_cache_ttl: <cache_ttl_in_seconds>
      langfuse_host: <langfuse_endpoint>
      langfuse_enable: <true/false>
-     enable_timer: <true/false>
      logging_level: <log_level>
      development: <true/false>
    ```

wren-ai-service/eval/evaluation.py

@@ -118,8 +118,8 @@ def _score_metrics(self, test_case: LLMTestCase, result: TestResult) -> None:
     @observe(name="Summary Trace", capture_input=False, capture_output=False)
     def _average_score(self, meta: dict) -> None:
         langfuse_context.update_current_trace(
-            session_id=meta["session_id"],
-            user_id=meta["user_id"],
+            session_id=meta.get("session_id"),
+            user_id=meta.get("user_id"),
             metadata=trace_metadata(meta, type="summary"),
         )
 

wren-ai-service/eval/pipelines.py

@@ -142,8 +142,8 @@ async def process(self, query: dict) -> dict:
         }
 
         langfuse_context.update_current_trace(
-            session_id=self._meta["session_id"],
-            user_id=self._meta["user_id"],
+            session_id=self._meta.get("session_id"),
+            user_id=self._meta.get("user_id"),
             metadata=trace_metadata(self._meta, type=prediction["type"]),
         )
 
@@ -159,8 +159,8 @@ async def flat(self, prediction: dict, **kwargs) -> dict:
 
         langfuse_context.update_current_trace(
             name=f"Prediction Process - Shallow Trace for {prediction['input']} ",
-            session_id=self._meta["session_id"],
-            user_id=self._meta["user_id"],
+            session_id=self._meta.get("session_id"),
+            user_id=self._meta.get("user_id"),
             metadata={
                 **trace_metadata(self._meta, type=prediction["type"]),
                 "source_trace_id": prediction["source_trace_id"],

wren-ai-service/src/config.py

@@ -46,7 +46,6 @@ class Settings(BaseSettings):
     langfuse_enable: bool = Field(default=True)
 
     # debug config
-    enable_timer: bool = Field(default=False)
     logging_level: str = Field(default="INFO")
     development: bool = Field(default=False)
 

wren-ai-service/src/globals.py

@@ -6,31 +6,7 @@
 from src.config import Settings
 from src.core.pipeline import PipelineComponent
 from src.core.provider import EmbedderProvider, LLMProvider
-from src.pipelines import indexing
-from src.pipelines.generation import (
-    chart_adjustment,
-    chart_generation,
-    data_assistance,
-    followup_sql_generation,
-    intent_classification,
-    question_recommendation,
-    relationship_recommendation,
-    semantics_description,
-    sql_answer,
-    sql_breakdown,
-    sql_correction,
-    sql_expansion,
-    sql_explanation,
-    sql_generation,
-    sql_regeneration,
-    sql_summary,
-)
-from src.pipelines.retrieval import (
-    historical_question,
-    preprocess_sql_data,
-    retrieval,
-    sql_executor,
-)
+from src.pipelines import generation, indexing, retrieval
 from src.web.v1.services.ask import AskService
 from src.web.v1.services.ask_details import AskDetailsService
 from src.web.v1.services.chart import ChartService
@@ -41,8 +17,9 @@
 from src.web.v1.services.semantics_preparation import SemanticsPreparationService
 from src.web.v1.services.sql_answer import SqlAnswerService
 from src.web.v1.services.sql_expansion import SqlExpansionService
-from src.web.v1.services.sql_explanation import SQLExplanationService
-from src.web.v1.services.sql_regeneration import SQLRegenerationService
+from src.web.v1.services.sql_explanation import SqlExplanationService
+from src.web.v1.services.sql_pairs_preparation import SqlPairsPreparationService
+from src.web.v1.services.sql_regeneration import SqlRegenerationService
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -59,8 +36,9 @@ class ServiceContainer:
     chart_adjustment_service: ChartAdjustmentService
     sql_answer_service: SqlAnswerService
     sql_expansion_service: SqlExpansionService
-    sql_explanation_service: SQLExplanationService
-    sql_regeneration_service: SQLRegenerationService
+    sql_explanation_service: SqlExplanationService
+    sql_regeneration_service: SqlRegenerationService
+    sql_pairs_preparation_service: SqlPairsPreparationService
 
 
 @dataclass
@@ -80,7 +58,7 @@ def create_service_container(
     return ServiceContainer(
         semantics_description=SemanticsDescription(
             pipelines={
-                "semantics_description": semantics_description.SemanticsDescription(
+                "semantics_description": generation.SemanticsDescription(
                     **pipe_components["semantics_description"],
                 )
             },
@@ -103,10 +81,10 @@ def create_service_container(
         ),
         ask_service=AskService(
             pipelines={
-                "intent_classification": intent_classification.IntentClassification(
+                "intent_classification": generation.IntentClassification(
                     **pipe_components["intent_classification"],
                 ),
-                "data_assistance": data_assistance.DataAssistance(
+                "data_assistance": generation.DataAssistance(
                     **pipe_components["data_assistance"]
                 ),
                 "retrieval": retrieval.Retrieval(
@@ -115,63 +93,66 @@ def create_service_container(
                     table_column_retrieval_size=settings.table_column_retrieval_size,
                     allow_using_db_schemas_without_pruning=settings.allow_using_db_schemas_without_pruning,
                 ),
-                "historical_question": historical_question.HistoricalQuestion(
+                "historical_question_retrieval": retrieval.HistoricalQuestionRetrieval(
                     **pipe_components["historical_question_retrieval"],
                 ),
-                "sql_generation": sql_generation.SQLGeneration(
+                "sql_pairs_retrieval": retrieval.SqlPairsRetrieval(
+                    **pipe_components["sql_pairs_retrieval"],
+                ),
+                "sql_generation": generation.SQLGeneration(
                     **pipe_components["sql_generation"],
                 ),
-                "sql_correction": sql_correction.SQLCorrection(
+                "sql_correction": generation.SQLCorrection(
                     **pipe_components["sql_correction"],
                 ),
-                "followup_sql_generation": followup_sql_generation.FollowUpSQLGeneration(
+                "followup_sql_generation": generation.FollowUpSQLGeneration(
                     **pipe_components["followup_sql_generation"],
                 ),
-                "sql_summary": sql_summary.SQLSummary(
+                "sql_summary": generation.SQLSummary(
                     **pipe_components["sql_summary"],
                 ),
             },
             **query_cache,
         ),
         chart_service=ChartService(
             pipelines={
-                "sql_executor": sql_executor.SQLExecutor(
+                "sql_executor": retrieval.SQLExecutor(
                     **pipe_components["sql_executor"],
                 ),
-                "chart_generation": chart_generation.ChartGeneration(
+                "chart_generation": generation.ChartGeneration(
                     **pipe_components["chart_generation"],
                 ),
             },
             **query_cache,
         ),
         chart_adjustment_service=ChartAdjustmentService(
             pipelines={
-                "sql_executor": sql_executor.SQLExecutor(
+                "sql_executor": retrieval.SQLExecutor(
                     **pipe_components["sql_executor"],
                 ),
-                "chart_adjustment": chart_adjustment.ChartAdjustment(
+                "chart_adjustment": generation.ChartAdjustment(
                     **pipe_components["chart_adjustment"],
                 ),
             },
             **query_cache,
         ),
         sql_answer_service=SqlAnswerService(
             pipelines={
-                "preprocess_sql_data": preprocess_sql_data.PreprocessSqlData(
+                "preprocess_sql_data": retrieval.PreprocessSqlData(
                     **pipe_components["preprocess_sql_data"],
                 ),
-                "sql_answer": sql_answer.SQLAnswer(
+                "sql_answer": generation.SQLAnswer(
                     **pipe_components["sql_answer"],
                 ),
             },
             **query_cache,
         ),
         ask_details_service=AskDetailsService(
             pipelines={
-                "sql_breakdown": sql_breakdown.SQLBreakdown(
+                "sql_breakdown": generation.SQLBreakdown(
                     **pipe_components["sql_breakdown"],
                 ),
-                "sql_summary": sql_summary.SQLSummary(
+                "sql_summary": generation.SQLSummary(
                     **pipe_components["sql_summary"],
                 ),
             },
@@ -184,45 +165,45 @@ def create_service_container(
                     table_retrieval_size=settings.table_retrieval_size,
                     table_column_retrieval_size=settings.table_column_retrieval_size,
                 ),
-                "sql_expansion": sql_expansion.SQLExpansion(
+                "sql_expansion": generation.SQLExpansion(
                     **pipe_components["sql_expansion"],
                 ),
-                "sql_correction": sql_correction.SQLCorrection(
+                "sql_correction": generation.SQLCorrection(
                     **pipe_components["sql_correction"],
                 ),
-                "sql_summary": sql_summary.SQLSummary(
+                "sql_summary": generation.SQLSummary(
                     **pipe_components["sql_summary"],
                 ),
             },
             **query_cache,
         ),
-        sql_explanation_service=SQLExplanationService(
+        sql_explanation_service=SqlExplanationService(
             pipelines={
-                "sql_explanation": sql_explanation.SQLExplanation(
+                "sql_explanation": generation.SQLExplanation(
                     **pipe_components["sql_explanation"],
                 )
             },
             **query_cache,
         ),
-        sql_regeneration_service=SQLRegenerationService(
+        sql_regeneration_service=SqlRegenerationService(
             pipelines={
-                "sql_regeneration": sql_regeneration.SQLRegeneration(
+                "sql_regeneration": generation.SQLRegeneration(
                     **pipe_components["sql_regeneration"],
                 )
             },
             **query_cache,
         ),
         relationship_recommendation=RelationshipRecommendation(
             pipelines={
-                "relationship_recommendation": relationship_recommendation.RelationshipRecommendation(
+                "relationship_recommendation": generation.RelationshipRecommendation(
                     **pipe_components["relationship_recommendation"],
                 )
             },
             **query_cache,
         ),
         question_recommendation=QuestionRecommendation(
             pipelines={
-                "question_recommendation": question_recommendation.QuestionRecommendation(
+                "question_recommendation": generation.QuestionRecommendation(
                     **pipe_components["question_recommendation"],
                 ),
                 "retrieval": retrieval.Retrieval(
@@ -231,12 +212,23 @@ def create_service_container(
                     table_column_retrieval_size=settings.table_column_retrieval_size,
                     allow_using_db_schemas_without_pruning=settings.allow_using_db_schemas_without_pruning,
                 ),
-                "sql_generation": sql_generation.SQLGeneration(
+                "sql_generation": generation.SQLGeneration(
                     **pipe_components["sql_generation"],
                 ),
             },
             **query_cache,
         ),
+        sql_pairs_preparation_service=SqlPairsPreparationService(
+            pipelines={
+                "sql_pairs_preparation": indexing.SqlPairsPreparation(
+                    **pipe_components["sql_pairs_preparation"],
+                ),
+                "sql_pairs_deletion": indexing.SqlPairsDeletion(
+                    **pipe_components["sql_pairs_deletion"],
+                ),
+            },
+            **query_cache,
+        ),
     )
 
 

wren-ai-service/src/pipelines/common.py

@@ -6,7 +6,7 @@
 import aiohttp
 import orjson
 import pytz
-from haystack import component
+from haystack import Document, component
 
 from src.core.engine import (
     Engine,
@@ -19,6 +19,21 @@
 logger = logging.getLogger("wren-ai-service")
 
 
+@component
+class ScoreFilter:
+    @component.output_types(
+        documents=List[Document],
+    )
+    def run(self, documents: List[Document], score: float = 0.9):
+        return {
+            "documents": sorted(
+                filter(lambda document: document.score >= score, documents),
+                key=lambda document: document.score,
+                reverse=True,
+            )
+        }
+
+
 @component
 class SQLBreakdownGenPostProcessor:
     def __init__(self, engine: Engine):