feat(wren-ai-service): user guide and misleading streaming(ai-env-changed) (#1015)

Author: cyyeh

Diff for: deployment/kustomizations/base/cm.yaml

@@ -154,6 +154,8 @@ data:
         llm: litellm_llm.default
         embedder: litellm_embedder.default
         document_store: qdrant
+      - name: misleading_assistance
+        llm: litellm_llm.default
       - name: data_assistance
         llm: litellm_llm.default
       - name: sql_pairs_indexing
@@ -171,6 +173,8 @@ data:
         llm: litellm_llm.default
       - name: chart_adjustment
         llm: litellm_llm.default
+      - name: user_guide_assistance
+        llm: litellm_llm.default
       - name: sql_question_generation
         llm: litellm_llm.default
       - name: sql_generation_reasoning
@@ -191,9 +195,10 @@ data:
         document_store: qdrant
       - name: project_meta_indexing
         document_store: qdrant
-
     ---
     settings:
+      doc_endpoint: https://docs.getwren.ai
+      is_oss: true
       engine_timeout: 30
       column_indexing_batch_size: 50
       table_retrieval_size: 10

Diff for: docker/config.example.yaml

@@ -106,6 +106,8 @@ pipes:
     llm: litellm_llm.default
     embedder: litellm_embedder.default
     document_store: qdrant
+  - name: misleading_assistance
+    llm: litellm_llm.default
   - name: data_assistance
     llm: litellm_llm.default
   - name: sql_pairs_indexing
@@ -123,6 +125,8 @@ pipes:
     llm: litellm_llm.default
   - name: chart_adjustment
     llm: litellm_llm.default
+  - name: user_guide_assistance
+    llm: litellm_llm.default
   - name: sql_question_generation
     llm: litellm_llm.default
   - name: sql_generation_reasoning
@@ -146,6 +150,8 @@ pipes:
 
 ---
 settings:
+  doc_endpoint: https://docs.getwren.ai
+  is_oss: true
   engine_timeout: 30
   column_indexing_batch_size: 50
   table_retrieval_size: 10

Diff for: wren-ai-service/docs/config_examples/config.anthropic.yaml

@@ -101,6 +101,8 @@ pipes:
     llm: litellm_llm.default
     embedder: litellm_embedder.default
     document_store: qdrant
+  - name: misleading_assistance
+    llm: litellm_llm.default
   - name: data_assistance
     llm: litellm_llm.default
   - name: sql_pairs_indexing
@@ -114,6 +116,8 @@ pipes:
     llm: litellm_llm.default
   - name: sql_executor
     engine: wren_ui
+  - name: user_guide_assistance
+    llm: litellm_llm.default
   - name: chart_generation
     llm: litellm_llm.default
   - name: chart_adjustment

Diff for: wren-ai-service/docs/config_examples/config.azure.yaml

@@ -104,6 +104,8 @@ pipes:
     llm: litellm_llm.default
     embedder: litellm_embedder.default
     document_store: qdrant
+  - name: misleading_assistance
+    llm: litellm_llm.default
   - name: data_assistance
     llm: litellm_llm.default
   - name: sql_pairs_preparation
@@ -118,6 +120,8 @@ pipes:
     llm: litellm_llm.default
   - name: sql_executor
     engine: wren_ui
+  - name: user_guide_assistance
+    llm: litellm_llm.default
   - name: chart_generation
     llm: litellm_llm.default
   - name: chart_adjustment

Diff for: wren-ai-service/docs/config_examples/config.deepseek.yaml

@@ -129,6 +129,8 @@ pipes:
     llm: litellm_llm.default
     embedder: litellm_embedder.default
     document_store: qdrant
+  - name: misleading_assistance
+    llm: litellm_llm.default
   - name: data_assistance
     llm: litellm_llm.deepseek/deepseek-chat
   - name: sql_pairs_indexing
@@ -142,6 +144,8 @@ pipes:
     llm: litellm_llm.default
   - name: sql_executor
     engine: wren_ui
+  - name: user_guide_assistance
+    llm: litellm_llm.default
   - name: sql_question_generation
     llm: litellm_llm.default
   - name: sql_generation_reasoning

Diff for: wren-ai-service/docs/config_examples/config.google_ai_studio.yaml

@@ -115,6 +115,8 @@ pipes:
     llm: litellm_llm.default
     embedder: litellm_embedder.default
     document_store: qdrant
+  - name: misleading_assistance
+    llm: litellm_llm.default
   - name: data_assistance
     llm: litellm_llm.default
   - name: sql_pairs_indexing
@@ -128,6 +130,8 @@ pipes:
     llm: litellm_llm.default
   - name: sql_executor
     engine: wren_ui
+  - name: user_guide_assistance
+    llm: litellm_llm.default
   - name: sql_question_generation
     llm: litellm_llm.default
   - name: sql_generation_reasoning

Diff for: wren-ai-service/docs/config_examples/config.groq.yaml

@@ -110,6 +110,8 @@ pipes:
     llm: litellm_llm.default
     embedder: litellm_embedder.default
     document_store: qdrant
+  - name: misleading_assistance
+    llm: litellm_llm.default
   - name: data_assistance
     llm: litellm_llm.default
   - name: sql_pairs_indexing
@@ -123,6 +125,8 @@ pipes:
     llm: litellm_llm.default
   - name: sql_executor
     engine: wren_ui
+  - name: user_guide_assistance
+    llm: litellm_llm.default
   - name: sql_question_generation
     llm: litellm_llm.default
   - name: sql_generation_reasoning

Diff for: wren-ai-service/docs/config_examples/config.lm_studio.yaml

@@ -109,6 +109,8 @@ pipes:
     llm: litellm_llm.default
     embedder: litellm_embedder.default
     document_store: qdrant
+  - name: misleading_assistance
+    llm: litellm_llm.default
   - name: data_assistance
     llm: litellm_llm.default
   - name: sql_pairs_indexing
@@ -122,6 +124,8 @@ pipes:
     llm: litellm_llm.default
   - name: sql_executor
     engine: wren_ui
+  - name: user_guide_assistance
+    llm: litellm_llm.default
   - name: sql_question_generation
     llm: litellm_llm.default
   - name: sql_generation_reasoning

Diff for: wren-ai-service/docs/config_examples/config.ollama.yaml

@@ -107,6 +107,8 @@ pipes:
     llm: litellm_llm.default
     embedder: litellm_embedder.default
     document_store: qdrant
+  - name: misleading_assistance
+    llm: litellm_llm.default
   - name: data_assistance
     llm: litellm_llm.default
   - name: sql_pairs_indexing
@@ -120,6 +122,8 @@ pipes:
     llm: litellm_llm.default
   - name: sql_executor
     engine: wren_ui
+  - name: user_guide_assistance
+    llm: litellm_llm.default
   - name: sql_question_generation
     llm: litellm_llm.default
   - name: sql_generation_reasoning

Diff for: wren-ai-service/src/config.py

@@ -54,6 +54,10 @@ class Settings(BaseSettings):
         """,
     )
 
+    # user guide config
+    is_oss: bool = Field(default=True)
+    doc_endpoint: str = Field(default="https://docs.getwren.ai")
+
     # langfuse config
     # in order to use langfuse, we also need to set the LANGFUSE_SECRET_KEY and LANGFUSE_PUBLIC_KEY in the .env or .env.dev file
     langfuse_host: str = Field(default="https://cloud.langfuse.com")

Diff for: wren-ai-service/src/globals.py

@@ -7,6 +7,7 @@
 from src.core.pipeline import PipelineComponent
 from src.core.provider import EmbedderProvider, LLMProvider
 from src.pipelines import generation, indexing, retrieval
+from src.utils import fetch_wren_ai_docs
 from src.web.v1 import services
 
 logger = logging.getLogger("wren-ai-service")
@@ -44,6 +45,10 @@ def create_service_container(
         "maxsize": settings.query_cache_maxsize,
         "ttl": settings.query_cache_ttl,
     }
+    wren_ai_docs = fetch_wren_ai_docs(settings.doc_endpoint, settings.is_oss)
+    if not wren_ai_docs:
+        logger.warning("Failed to fetch Wren AI docs or response was empty.")
+
     return ServiceContainer(
         semantics_description=services.SemanticsDescription(
             pipelines={
@@ -82,10 +87,18 @@ def create_service_container(
             pipelines={
                 "intent_classification": generation.IntentClassification(
                     **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,
+                ),
                 "retrieval": retrieval.Retrieval(
                     **pipe_components["db_schema_retrieval"],
                     table_retrieval_size=settings.table_retrieval_size,

Diff for: wren-ai-service/src/pipelines/generation/__init__.py

@@ -4,6 +4,7 @@
 from .followup_sql_generation import FollowUpSQLGeneration
 from .followup_sql_generation_reasoning import FollowUpSQLGenerationReasoning
 from .intent_classification import IntentClassification
+from .misleading_assistance import MisleadingAssistance
 from .question_recommendation import QuestionRecommendation
 from .relationship_recommendation import RelationshipRecommendation
 from .semantics_description import SemanticsDescription
@@ -16,6 +17,7 @@
 from .sql_question import SQLQuestion
 from .sql_regeneration import SQLRegeneration
 from .sql_summary import SQLSummary
+from .user_guide_assistance import UserGuideAssistance
 
 __all__ = [
     "ChartGeneration",
@@ -33,7 +35,9 @@
     "SQLGeneration",
     "SQLGenerationReasoning",
     "SQLSummary",
+    "UserGuideAssistance",
     "SQLQuestion",
     "SQLRegeneration",
     "FollowUpSQLGenerationReasoning",
+    "MisleadingAssistance",
 ]

Diff for: wren-ai-service/src/pipelines/generation/intent_classification.py

@@ -25,7 +25,7 @@
 ### TASK ###
 You are a great detective, who is great at intent classification.
 First, rephrase the user's question to make it more specific, clear and relevant to the database schema before making the intent classification.
-Second, you need to use rephrased user's question to classify user's intent based on given database schema to one of three conditions: MISLEADING_QUERY, TEXT_TO_SQL, GENERAL. 
+Second, you need to use rephrased user's question to classify user's intent based on given database schema to one of four conditions: MISLEADING_QUERY, TEXT_TO_SQL, GENERAL, USER_GUIDE. 
 Also you should provide reasoning for the classification clearly and concisely within 20 words.
 
 ### INSTRUCTIONS ###
@@ -56,6 +56,31 @@
         - "What is the total sales for last quarter?"
         - "Show me all customers who purchased product X."
         - "List the top 10 products by revenue."
+- GENERAL
+    - When to Use:
+        - Use this category if the user is seeking general information about the database schema.
+        - If the rephrasedd user's question is related to the previous question, but considering them together cannot be answered by generating an SQL query using that schema.
+    - Characteristics:
+        - The question is about understanding the dataset or its capabilities.
+        - The user may need guidance on how to proceed or what questions to ask.
+    - Instructions:
+        - MUST explicitly add phrases from the rephrasedd user's question that are not explicitly related to the database schema in the reasoning output. Choose the most relevant phrases that cause the rephrasedd user's question to be GENERAL.
+    - Examples:
+        - "What is the dataset about?"
+        - "Tell me more about the database."
+        - "How can I analyze customer behavior with this data?"
+- USER_GUIDE
+    - When to Use:
+        - If the user's question is about Wren AI's features, capabilities, or how to use Wren AI.
+        - If the user's question is related to the content in the user guide.
+    - Characteristics:
+        - The question is about Wren AI's features, capabilities, or how to use Wren AI.
+    - Examples:
+        - "What can Wren AI do?"
+        - "How can I reset project?"
+        - "How can I delete project?"
+        - "How can I connect to other databases?"
+        - "How to draw a chart?"
 - MISLEADING_QUERY
     - When to Use:
         - If the rephrasedd user's question is irrelevant to the given database schema and cannot be answered using SQL with that schema.
@@ -71,28 +96,14 @@
         - "How are you?"
         - "What's the weather like today?"
         - "Tell me a joke."
-- GENERAL
-    - When to Use:
-        - Use this category if the user is seeking general information about the database schema.
-        - If the rephrasedd user's question is related to the previous question, but considering them together cannot be answered by generating an SQL query using that schema.
-    - Characteristics:
-        - The question is about understanding the dataset or its capabilities.
-        - The user may need guidance on how to proceed or what questions to ask.
-    - Instructions:
-        - MUST explicitly add phrases from the rephrasedd user's question that are not explicitly related to the database schema in the reasoning output. Choose the most relevant phrases that cause the rephrasedd user's question to be GENERAL.
-    - Examples:
-        - "What is the dataset about?"
-        - "Tell me more about the database."
-        - "What can Wren AI do?"
-        - "How can I analyze customer behavior with this data?"
-
+        
 ### OUTPUT FORMAT ###
 Please provide your response as a JSON object, structured as follows:
 
 {
     "rephrased_question": "<REPHRASED_USER_QUESTION_IN_STRING_FORMAT>",
     "reasoning": "<CHAIN_OF_THOUGHT_REASONING_BASED_ON_REPHRASED_USER_QUESTION_IN_STRING_FORMAT>",
-    "results": "MISLEADING_QUERY" | "TEXT_TO_SQL" | "GENERAL"
+    "results": "MISLEADING_QUERY" | "TEXT_TO_SQL" | "GENERAL" | "USER_GUIDE"
 }
 """
 
@@ -127,6 +138,11 @@
 {% endfor %}
 {% endif %}
 
+### USER GUIDE ###
+{% for doc in docs %}
+- {{doc.path}}: {{doc.content}}
+{% endfor %}
+
 ### QUESTION ###
 User's question: {{query}}
 Current Time: {{ current_time }}
@@ -245,6 +261,7 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[str]:
 @observe(capture_input=False)
 def prompt(
     query: str,
+    wren_ai_docs: list[dict],
     construct_db_schemas: list[str],
     prompt_builder: PromptBuilder,
     histories: Optional[list[AskHistory]] = None,
@@ -263,6 +280,7 @@ def prompt(
             configuration=configuration,
         ),
         current_time=configuration.show_current_time(),
+        docs=wren_ai_docs,
     )
 
 
@@ -294,7 +312,7 @@ def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict
 
 
 class IntentClassificationResult(BaseModel):
-    results: Literal["MISLEADING_QUERY", "TEXT_TO_SQL", "GENERAL"]
+    results: Literal["MISLEADING_QUERY", "TEXT_TO_SQL", "GENERAL", "USER_GUIDE"]
     rephrased_question: str
     reasoning: str
 
@@ -316,6 +334,7 @@ def __init__(
         llm_provider: LLMProvider,
         embedder_provider: EmbedderProvider,
         document_store_provider: DocumentStoreProvider,
+        wren_ai_docs: list[dict],
         table_retrieval_size: Optional[int] = 50,
         table_column_retrieval_size: Optional[int] = 100,
         **kwargs,
@@ -339,6 +358,10 @@ def __init__(
             ),
         }
 
+        self._configs = {
+            "wren_ai_docs": wren_ai_docs,
+        }
+
         super().__init__(
             AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
         )
@@ -364,6 +387,7 @@ async def run(
                 "instructions": instructions or [],
                 "configuration": configuration,
                 **self._components,
+                **self._configs,
             },
         )