feat(wren-ai-service): new-workflow-api (ai-env-changed)

deployment/kustomizations/base/cm.yaml

@@ -189,7 +189,8 @@ data:
         document_store: qdrant
       - name: sql_tables_extraction
         llm: litellm_llm.default
-
+      - name: data_exploration_assistance
+        llm: litellm_llm.default
     ---
     settings:
       doc_endpoint: https://docs.getwren.ai

docker/config.example.yaml

@@ -139,7 +139,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   doc_endpoint: https://docs.getwren.ai

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

@@ -136,7 +136,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

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

@@ -149,7 +149,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

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

@@ -159,7 +159,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

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

@@ -145,7 +145,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

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

@@ -141,7 +141,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

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

@@ -140,7 +140,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

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

@@ -139,7 +139,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

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

@@ -137,7 +137,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

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

@@ -137,7 +137,8 @@ pipes:
     document_store: qdrant
   - name: sql_tables_extraction
     llm: litellm_llm.default
-
+  - name: data_exploration_assistance
+    llm: litellm_llm.default
 ---
 settings:
   engine_timeout: 30

wren-ai-service/pyproject.toml

@@ -39,26 +39,19 @@ qdrant-client = "==1.11.0"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.7.1"
-streamlit = "^1.37.0"
-watchdog = "^4.0.0"
 pandas = "^2.2.2"
-matplotlib = "^3.9.2"
-sseclient-py = "^1.8.0"
-dspy-ai = "^2.5.26"
 requests = "^2.32.2"
-extra-streamlit-components = "^0.1.71"
-deepeval = "^1.0.6"
-tomlkit = "^0.13.0"
-nltk = "^3.9.1"
 
 [tool.poetry.group.eval.dependencies]
-gitpython = "^3.1.43"
 plotly = "^5.24.1"
 nbformat = "^5.1.3"
 ipykernel = "^6.29.5"
 itables = "^2.2.1"
 gdown = "^5.2.0"
-streamlit-tags = "^1.2.8"
+nltk = "^3.9.1"
+tomlkit = "^0.13.0"
+deepeval = "^1.0.6"
+dspy-ai = "^2.5.26"
 
 [tool.poetry.group.test.dependencies]
 locust = "^2.32.0"

wren-ai-service/src/__main__.py

@@ -18,6 +18,7 @@
     setup_custom_logger,
 )
 from src.web.v1 import routers
+from src.web.v2 import routers as v2_routers
 
 setup_custom_logger(
     "wren-ai-service", level_str=settings.logging_level, is_dev=settings.development
@@ -54,6 +55,7 @@ async def lifespan(app: FastAPI):
     allow_headers=["*"],
 )
 app.include_router(routers.router, prefix="/v1", tags=["v1"])
+app.include_router(v2_routers.router, prefix="/v2", tags=["v2"])
 # TODO: deprecated, it was used for load testing using locust only. should be removed in the future
 if settings.development:
     from src.web import development

wren-ai-service/src/globals.py

@@ -9,6 +9,7 @@
 from src.pipelines import generation, indexing, retrieval
 from src.utils import fetch_wren_ai_docs
 from src.web.v1 import services
+from src.web.v2 import services as v2_services
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -27,6 +28,7 @@ class ServiceContainer:
     sql_question_service: services.SqlQuestionService
     instructions_service: services.InstructionsService
     sql_correction_service: services.SqlCorrectionService
+    conversation_service: v2_services.ConversationService
 
 
 @dataclass
@@ -269,6 +271,87 @@ def create_service_container(
             },
             **query_cache,
         ),
+        conversation_service=v2_services.ConversationService(
+            pipelines={
+                "intent_classification": generation.IntentClassificationV2(
+                    **pipe_components["intent_classification"],
+                    wren_ai_docs=wren_ai_docs,
+                ),
+                "misleading_assistance": generation.MisleadingAssistance(
+                    **pipe_components["misleading_assistance"],
+                ),
+                "data_assistance": generation.DataAssistance(
+                    **pipe_components["data_assistance"]
+                ),
+                "user_guide_assistance": generation.UserGuideAssistance(
+                    **pipe_components["user_guide_assistance"],
+                    wren_ai_docs=wren_ai_docs,
+                ),
+                "data_exploration_assistance": generation.DataExplorationAssistance(
+                    **pipe_components["data_exploration_assistance"],
+                ),
+                "db_schema_retrieval": retrieval.DbSchemaRetrieval(
+                    **pipe_components["db_schema_retrieval"],
+                    table_retrieval_size=settings.table_retrieval_size,
+                    table_column_retrieval_size=settings.table_column_retrieval_size,
+                ),
+                "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"],
+                    sql_pairs_similarity_threshold=settings.sql_pairs_similarity_threshold,
+                    sql_pairs_retrieval_max_size=settings.sql_pairs_retrieval_max_size,
+                ),
+                "instructions_retrieval": retrieval.Instructions(
+                    **pipe_components["instructions_retrieval"],
+                    similarity_threshold=settings.instructions_similarity_threshold,
+                    top_k=settings.instructions_top_k,
+                ),
+                "sql_generation": generation.SQLGeneration(
+                    **pipe_components["sql_generation"],
+                    engine_timeout=settings.engine_timeout,
+                ),
+                "sql_generation_reasoning": generation.SQLGenerationReasoning(
+                    **pipe_components["sql_generation_reasoning"],
+                ),
+                "followup_sql_generation_reasoning": generation.FollowUpSQLGenerationReasoning(
+                    **pipe_components["followup_sql_generation_reasoning"],
+                ),
+                "sql_correction": generation.SQLCorrection(
+                    **pipe_components["sql_correction"],
+                    engine_timeout=settings.engine_timeout,
+                ),
+                "followup_sql_generation": generation.FollowUpSQLGeneration(
+                    **pipe_components["followup_sql_generation"],
+                    engine_timeout=settings.engine_timeout,
+                ),
+                "sql_regeneration": generation.SQLRegeneration(
+                    **pipe_components["sql_regeneration"],
+                    engine_timeout=settings.engine_timeout,
+                ),
+                "sql_functions_retrieval": retrieval.SqlFunctions(
+                    **pipe_components["sql_functions_retrieval"],
+                    engine_timeout=settings.engine_timeout,
+                ),
+                "sql_executor": retrieval.SQLExecutor(
+                    **pipe_components["sql_executor"],
+                    engine_timeout=settings.engine_timeout,
+                ),
+                "sql_answer": generation.SQLAnswer(
+                    **pipe_components["sql_answer"],
+                    engine_timeout=settings.engine_timeout,
+                ),
+                "chart_generation": generation.ChartGeneration(
+                    **pipe_components["chart_generation"],
+                ),
+                "chart_adjustment": generation.ChartAdjustmentV2(
+                    **pipe_components["chart_adjustment"],
+                ),
+            },
+            max_histories=settings.max_histories,
+        ),
     )
 
 

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

@@ -1,9 +1,12 @@
 from .chart_adjustment import ChartAdjustment
+from .chart_adjustment_v2 import ChartAdjustmentV2
 from .chart_generation import ChartGeneration
 from .data_assistance import DataAssistance
+from .data_exploration_assistance import DataExplorationAssistance
 from .followup_sql_generation import FollowUpSQLGeneration
 from .followup_sql_generation_reasoning import FollowUpSQLGenerationReasoning
 from .intent_classification import IntentClassification
+from .intent_classification_v2 import IntentClassificationV2
 from .misleading_assistance import MisleadingAssistance
 from .question_recommendation import QuestionRecommendation
 from .relationship_recommendation import RelationshipRecommendation
@@ -20,9 +23,11 @@
 __all__ = [
     "ChartGeneration",
     "ChartAdjustment",
+    "ChartAdjustmentV2",
     "DataAssistance",
     "FollowUpSQLGeneration",
     "IntentClassification",
+    "IntentClassificationV2",
     "QuestionRecommendation",
     "RelationshipRecommendation",
     "SemanticsDescription",
@@ -36,4 +41,5 @@
     "FollowUpSQLGenerationReasoning",
     "MisleadingAssistance",
     "SQLTablesExtraction",
+    "DataExplorationAssistance",
 ]