update

Author: cyyeh

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

@@ -41,6 +41,13 @@
     {{ db_schema }}
 {% endfor %}
 
+{% if histories %}
+### PREVIOUS QUESTIONS ###
+{% for history in histories %}
+    {{ history.question }}
+{% endfor %}
+{% endif %}
+
 ### INPUT ###
 User's question: {{query}}
 Language: {{language}}
@@ -61,13 +68,9 @@ def prompt(
     prompt_builder: PromptBuilder,
     custom_instruction: str,
 ) -> dict:
-    previous_query_summaries = (
-        [history.question for history in histories] if histories else []
-    )
-    query = "\n".join(previous_query_summaries) + "\n" + query
-
     _prompt = prompt_builder.run(
         query=query,
+        histories=histories,
         db_schemas=db_schemas,
         language=language,
         custom_instruction=custom_instruction,

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

@@ -10,6 +10,9 @@
 
 from src.core.pipeline import BasicPipeline
 from src.core.provider import LLMProvider
+from src.pipelines.common import clean_up_new_lines
+from src.utils import trace_cost
+from src.web.v1.services.ask import AskHistory
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -32,10 +35,21 @@
 """
 
 data_exploration_assistance_user_prompt_template = """
+{% if histories %}
+### PREVIOUS QUESTIONS ###
+{% for history in histories %}
+    {{ history.question }}
+{% endfor %}
+{% endif %}
+
+### INPUT ###
 User Question: {{query}}
 Language: {{language}}
 SQL Data:
 {{ sql_data }}
+
+Custom Instruction: {{ custom_instruction }}
+
 Please think step by step.
 """
 
@@ -44,18 +58,24 @@
 @observe(capture_input=False)
 def prompt(
     query: str,
+    histories: list[AskHistory],
     language: str,
     sql_data: dict,
     prompt_builder: PromptBuilder,
+    custom_instruction: str,
 ) -> dict:
-    return prompt_builder.run(
+    _prompt = prompt_builder.run(
         query=query,
         language=language,
         sql_data=sql_data,
+        histories=histories,
+        custom_instruction=custom_instruction,
     )
+    return {"prompt": clean_up_new_lines(_prompt.get("prompt"))}
 
 
 @observe(as_type="generation", capture_input=False)
+@trace_cost
 async def data_exploration_assistance(
     prompt: dict, generator: Any, query_id: str
 ) -> dict:
@@ -127,6 +147,8 @@ async def run(
         sql_data: dict,
         language: str,
         query_id: Optional[str] = None,
+        histories: Optional[list[AskHistory]] = None,
+        custom_instruction: Optional[str] = None,
     ):
         logger.info("Data Exploration Assistance pipeline is running...")
         return await self._pipe.execute(
@@ -136,6 +158,8 @@ async def run(
                 "language": language,
                 "query_id": query_id or "",
                 "sql_data": sql_data,
+                "histories": histories or [],
+                "custom_instruction": custom_instruction or "",
                 **self._components,
             },
         )

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

@@ -41,6 +41,13 @@
     {{ db_schema }}
 {% endfor %}
 
+{% if histories %}
+### PREVIOUS QUESTIONS ###
+{% for history in histories %}
+    {{ history.question }}
+{% endfor %}
+{% endif %}
+
 ### INPUT ###
 User's question: {{query}}
 Language: {{language}}
@@ -61,13 +68,9 @@ def prompt(
     prompt_builder: PromptBuilder,
     custom_instruction: str,
 ) -> dict:
-    previous_query_summaries = (
-        [history.question for history in histories] if histories else []
-    )
-    query = "\n".join(previous_query_summaries) + "\n" + query
-
     _prompt = prompt_builder.run(
         query=query,
+        histories=histories,
         db_schemas=db_schemas,
         language=language,
         custom_instruction=custom_instruction,

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

@@ -34,13 +34,9 @@ def prompt(
     prompt_builder: PromptBuilder,
     custom_instruction: str,
 ) -> dict:
-    previous_query_summaries = (
-        [history.question for history in histories] if histories else []
-    )
-    query = "\n".join(previous_query_summaries) + "\n" + query
-
     _prompt = prompt_builder.run(
         query=query,
+        histories=histories,
         db_schemas=db_schemas,
         language=language,
         custom_instruction=custom_instruction,

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

@@ -12,6 +12,7 @@
 from src.core.provider import LLMProvider
 from src.pipelines.common import clean_up_new_lines
 from src.utils import trace_cost
+from src.web.v1.services.ask import AskHistory
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -34,6 +35,14 @@
 """
 
 user_guide_assistance_user_prompt_template = """
+{% if histories %}
+### PREVIOUS QUESTIONS ###
+{% for history in histories %}
+    {{ history.question }}
+{% endfor %}
+{% endif %}
+
+### INPUT ###
 User Question: {{query}}
 Language: {{language}}
 User Guide:
@@ -53,11 +62,13 @@ def prompt(
     query: str,
     language: str,
     wren_ai_docs: list[dict],
+    histories: list[AskHistory],
     prompt_builder: PromptBuilder,
     custom_instruction: str,
 ) -> dict:
     _prompt = prompt_builder.run(
         query=query,
+        histories=histories,
         language=language,
         docs=wren_ai_docs,
         custom_instruction=custom_instruction,
@@ -144,6 +155,7 @@ async def run(
         query: str,
         language: str,
         query_id: Optional[str] = None,
+        histories: Optional[list[AskHistory]] = None,
         custom_instruction: Optional[str] = None,
     ):
         logger.info("User Guide Assistance pipeline is running...")
@@ -153,6 +165,7 @@ async def run(
                 "query": query,
                 "language": language,
                 "query_id": query_id or "",
+                "histories": histories or [],
                 "custom_instruction": custom_instruction or "",
                 **self._components,
                 **self._configs,

wren-ai-service/src/web/v1/services/ask.py

@@ -336,6 +336,7 @@ async def ask(
                             asyncio.create_task(
                                 self._pipelines["user_guide_assistance"].run(
                                     query=user_query,
+                                    histories=histories,
                                     language=ask_request.configurations.language,
                                     query_id=ask_request.query_id,
                                     custom_instruction=ask_request.custom_instruction,
@@ -357,6 +358,7 @@ async def ask(
                             asyncio.create_task(
                                 self._pipelines["data_exploration_assistance"].run(
                                     query=user_query,
+                                    histories=histories,
                                     sql_data=last_sql_data,
                                     language=ask_request.configurations.language,
                                     query_id=ask_request.query_id,
@@ -379,6 +381,7 @@ async def ask(
                             asyncio.create_task(
                                 self._pipelines["user_clarification_assistance"].run(
                                     query=user_query,
+                                    histories=histories,
                                     db_schemas=table_ddls,
                                     language=ask_request.configurations.language,
                                     query_id=ask_request.query_id,