From e1fd1644d529b4420ad77aee7aa5fb8267f2367a Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 15:17:40 +0800 Subject: [PATCH 01/30] chore: update the all config for instructions indexing --- .../docs/config_examples/config.azure.yaml | 3 ++ .../docs/config_examples/config.deepseek.yaml | 3 ++ .../config.google_ai_studio.yaml | 3 ++ .../docs/config_examples/config.groq.yaml | 3 ++ .../docs/config_examples/config.ollama.yaml | 3 ++ .../tools/config/config.example.yaml | 51 ++++++++++--------- wren-ai-service/tools/config/config.full.yaml | 3 ++ 7 files changed, 45 insertions(+), 24 deletions(-) diff --git a/wren-ai-service/docs/config_examples/config.azure.yaml b/wren-ai-service/docs/config_examples/config.azure.yaml index 3785ae0c3a..2e219bf080 100644 --- a/wren-ai-service/docs/config_examples/config.azure.yaml +++ b/wren-ai-service/docs/config_examples/config.azure.yaml @@ -132,6 +132,9 @@ pipes: - name: sql_regeneration llm: litellm_llm.default engine: wren_ui + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant --- settings: diff --git a/wren-ai-service/docs/config_examples/config.deepseek.yaml b/wren-ai-service/docs/config_examples/config.deepseek.yaml index a8c20b3154..0a657e2b87 100644 --- a/wren-ai-service/docs/config_examples/config.deepseek.yaml +++ b/wren-ai-service/docs/config_examples/config.deepseek.yaml @@ -142,6 +142,9 @@ pipes: - name: sql_regeneration llm: litellm_llm.default engine: wren_ui + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant --- settings: diff --git a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml index 52cc6534df..0a5e7434aa 100644 --- a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml +++ b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml @@ -131,6 +131,9 @@ pipes: - name: sql_regeneration llm: litellm_llm.default engine: wren_ui + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant --- settings: diff --git a/wren-ai-service/docs/config_examples/config.groq.yaml b/wren-ai-service/docs/config_examples/config.groq.yaml index 780b0eadf3..a9895a5d42 100644 --- a/wren-ai-service/docs/config_examples/config.groq.yaml +++ b/wren-ai-service/docs/config_examples/config.groq.yaml @@ -124,6 +124,9 @@ pipes: - name: sql_regeneration llm: litellm_llm.default engine: wren_ui + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant --- settings: diff --git a/wren-ai-service/docs/config_examples/config.ollama.yaml b/wren-ai-service/docs/config_examples/config.ollama.yaml index 5beef2ef2a..8a60acd40c 100644 --- a/wren-ai-service/docs/config_examples/config.ollama.yaml +++ b/wren-ai-service/docs/config_examples/config.ollama.yaml @@ -121,6 +121,9 @@ pipes: - name: sql_regeneration llm: litellm_llm.default engine: wren_ui + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant --- settings: diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index 7f4adbabb8..1c3af50a4c 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -2,30 +2,30 @@ type: llm provider: litellm_llm timeout: 120 models: -- model: gpt-4o-mini-2024-07-18 - alias: default - api_base: https://api.openai.com/v1 - kwargs: - temperature: 0 - n: 1 - # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed - seed: 0 - max_tokens: 4096 -- model: gpt-4o-2024-08-06 - api_base: https://api.openai.com/v1 - kwargs: - temperature: 0 - n: 1 - # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed - seed: 0 - max_tokens: 4096 -- model: o3-mini-2025-01-31 - api_base: https://api.openai.com/v1 - kwargs: - n: 1 - seed: 0 - max_completion_tokens: 4096 - reasoning_effort: low + - model: gpt-4o-mini-2024-07-18 + alias: default + api_base: https://api.openai.com/v1 + kwargs: + temperature: 0 + n: 1 + # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed + seed: 0 + max_tokens: 4096 + - model: gpt-4o-2024-08-06 + api_base: https://api.openai.com/v1 + kwargs: + temperature: 0 + n: 1 + # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed + seed: 0 + max_tokens: 4096 + - model: o3-mini-2025-01-31 + api_base: https://api.openai.com/v1 + kwargs: + n: 1 + seed: 0 + max_completion_tokens: 4096 + reasoning_effort: low --- type: embedder @@ -145,6 +145,9 @@ pipes: engine: wren_ui - name: evaluation llm: litellm_llm.default + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant --- settings: diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 6e79eebdf1..49c27f02b0 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -145,6 +145,9 @@ pipes: engine: wren_ui - name: evaluation llm: litellm_llm.default + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant --- settings: From e5370b1ebd52d2ec08db48f68b43ab8eba59210c Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 15:33:03 +0800 Subject: [PATCH 02/30] feat: instructions indexing and clean pipelien --- .../src/pipelines/indexing/instructions.py | 191 ++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 wren-ai-service/src/pipelines/indexing/instructions.py diff --git a/wren-ai-service/src/pipelines/indexing/instructions.py b/wren-ai-service/src/pipelines/indexing/instructions.py new file mode 100644 index 0000000000..2a8581777d --- /dev/null +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -0,0 +1,191 @@ +import logging +import sys +import uuid +from typing import Any, Dict, List, Optional + +from hamilton import base +from hamilton.async_driver import AsyncDriver +from haystack import Document, component +from haystack.document_stores.types import DocumentStore, DuplicatePolicy +from langfuse.decorators import observe +from pydantic import BaseModel + +from src.core.pipeline import BasicPipeline +from src.core.provider import DocumentStoreProvider, EmbedderProvider +from src.pipelines.indexing import AsyncDocumentWriter + +logger = logging.getLogger("wren-ai-service") + + +class Instruction(BaseModel): + id: str + instruction: str = "" + question: str = "" + + +@component +class InstructionsConverter: + @component.output_types(documents=List[Document]) + def run(self, instructions: list[Instruction], project_id: Optional[str] = ""): + logger.info(f"Project ID: {project_id} Converting instructions to documents...") + + addition = {"project_id": project_id} if project_id else {} + + # todo: for the prototype, the content using instruction directly + return { + "documents": [ + Document( + id=str(uuid.uuid4()), + meta={ + "instruction_id": instruction.id, + "question": instruction.question, + **addition, + }, + content=instruction.instruction, + ) + for instruction in instructions + ] + } + + +@component +class InstructionsCleaner: + def __init__(self, instructions_store: DocumentStore) -> None: + self.store = instructions_store + + @component.output_types() + async def run( + self, instruction_ids: List[str], project_id: Optional[str] = None + ) -> None: + filter = { + "operator": "AND", + "conditions": [ + {"field": "instruction_id", "operator": "in", "value": instruction_ids}, + ], + } + + if project_id: + filter["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + + return await self.store.delete_documents(filter) + + +## Start of Pipeline + + +@observe(capture_input=False) +def to_documents( + instructions: List[Instruction], + document_converter: InstructionsConverter, + project_id: Optional[str] = "", +) -> Dict[str, Any]: + return document_converter.run(instructions=instructions, project_id=project_id) + + +@observe(capture_input=False, capture_output=False) +async def embedding( + to_documents: Dict[str, Any], + embedder: Any, +) -> Dict[str, Any]: + return await embedder.run(documents=to_documents["documents"]) + + +@observe(capture_input=False, capture_output=False) +async def clean( + cleaner: InstructionsCleaner, + instructions: List[Instruction], + embedding: Dict[str, Any], + project_id: Optional[str] = "", + delete_all: bool = False, +) -> Dict[str, Any]: + instruction_ids = [instruction.id for instruction in instructions] + if instruction_ids or delete_all: + await cleaner.run(instruction_ids=instruction_ids, project_id=project_id) + + return embedding + + +@observe(capture_input=False) +async def write( + clean: Dict[str, Any], + writer: AsyncDocumentWriter, +) -> None: + return await writer.run(documents=clean["documents"]) + + +## End of Pipeline + + +class Instructions(BasicPipeline): + def __init__( + self, + embedder_provider: EmbedderProvider, + document_store_provider: DocumentStoreProvider, + **kwargs, + ) -> None: + store = document_store_provider.get_store(dataset_name="instructions") + + self._components = { + "cleaner": InstructionsCleaner(store), + "embedder": embedder_provider.get_document_embedder(), + "document_converter": InstructionsConverter(), + "writer": AsyncDocumentWriter( + document_store=store, + policy=DuplicatePolicy.OVERWRITE, + ), + } + + super().__init__( + AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult()) + ) + + @observe(name="Instructions Indexing") + async def run( + self, + project_id: Optional[str] = "", + instructions: list[Instruction] = [], + ) -> Dict[str, Any]: + logger.info( + f"Project ID: {project_id} Instructions Indexing pipeline is running..." + ) + + input = { + "project_id": project_id, + "instructions": instructions, + **self._components, + } + + return await self._pipe.execute(["write"], inputs=input) + + @observe(name="Clean Documents for Instructions") + async def clean( + self, + instructions: List[Instruction], + project_id: Optional[str] = None, + delete_all: bool = False, + ) -> None: + await clean( + instructions=instructions, + embedding={"documents": []}, + cleaner=self._components["cleaner"], + project_id=project_id, + delete_all=delete_all, + ) + + +if __name__ == "__main__": + from src.pipelines.common import dry_run_pipeline + + dry_run_pipeline( + Instructions, + "instructions_indexing", + instructions=[ + Instruction( + id="1", + instruction="France is in the table of customers and the column is country", + question="What is the capital of France?", + ) + ], + ) From 0b2ee3941cd3ba2ea7369168ebaa95f1b204d159 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 15:35:34 +0800 Subject: [PATCH 03/30] chore: update all config for instructions retrieval --- wren-ai-service/docs/config_examples/config.azure.yaml | 3 +++ wren-ai-service/docs/config_examples/config.deepseek.yaml | 3 +++ .../docs/config_examples/config.google_ai_studio.yaml | 3 +++ wren-ai-service/docs/config_examples/config.groq.yaml | 3 +++ wren-ai-service/docs/config_examples/config.ollama.yaml | 3 +++ wren-ai-service/tools/config/config.example.yaml | 3 +++ wren-ai-service/tools/config/config.full.yaml | 3 +++ 7 files changed, 21 insertions(+) diff --git a/wren-ai-service/docs/config_examples/config.azure.yaml b/wren-ai-service/docs/config_examples/config.azure.yaml index 2e219bf080..a8f2da1186 100644 --- a/wren-ai-service/docs/config_examples/config.azure.yaml +++ b/wren-ai-service/docs/config_examples/config.azure.yaml @@ -135,6 +135,9 @@ pipes: - name: instructions_indexing embedder: litellm_embedder.default document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.azure/text-embedding-ada-002 + document_store: qdrant --- settings: diff --git a/wren-ai-service/docs/config_examples/config.deepseek.yaml b/wren-ai-service/docs/config_examples/config.deepseek.yaml index 0a657e2b87..5b18bd2cc6 100644 --- a/wren-ai-service/docs/config_examples/config.deepseek.yaml +++ b/wren-ai-service/docs/config_examples/config.deepseek.yaml @@ -145,6 +145,9 @@ pipes: - name: instructions_indexing embedder: litellm_embedder.default document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.text-embedding-3-large + document_store: qdrant --- settings: diff --git a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml index 0a5e7434aa..680f07942e 100644 --- a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml +++ b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml @@ -134,6 +134,9 @@ pipes: - name: instructions_indexing embedder: litellm_embedder.default document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.gemini/text-embedding-004 + document_store: qdrant --- settings: diff --git a/wren-ai-service/docs/config_examples/config.groq.yaml b/wren-ai-service/docs/config_examples/config.groq.yaml index a9895a5d42..d4a8fe915f 100644 --- a/wren-ai-service/docs/config_examples/config.groq.yaml +++ b/wren-ai-service/docs/config_examples/config.groq.yaml @@ -127,6 +127,9 @@ pipes: - name: instructions_indexing embedder: litellm_embedder.default document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.text-embedding-3-large + document_store: qdrant --- settings: diff --git a/wren-ai-service/docs/config_examples/config.ollama.yaml b/wren-ai-service/docs/config_examples/config.ollama.yaml index 8a60acd40c..26c9faf53c 100644 --- a/wren-ai-service/docs/config_examples/config.ollama.yaml +++ b/wren-ai-service/docs/config_examples/config.ollama.yaml @@ -124,6 +124,9 @@ pipes: - name: instructions_indexing embedder: litellm_embedder.default document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.openai/nomic-embed-text + document_store: qdrant --- settings: diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index 1c3af50a4c..e4b2a5914d 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -148,6 +148,9 @@ pipes: - name: instructions_indexing embedder: litellm_embedder.default document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.text-embedding-3-large + document_store: qdrant --- settings: diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 49c27f02b0..001f43ea34 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -148,6 +148,9 @@ pipes: - name: instructions_indexing embedder: litellm_embedder.default document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.text-embedding-3-large + document_store: qdrant --- settings: From 04d09d5fb1659ee975ed49421a155e4c33cce9a9 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 15:36:14 +0800 Subject: [PATCH 04/30] feat: implement instruction retrieval pipeline --- .../src/pipelines/retrieval/instructions.py | 168 ++++++++++++++++++ 1 file changed, 168 insertions(+) create mode 100644 wren-ai-service/src/pipelines/retrieval/instructions.py diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py new file mode 100644 index 0000000000..5bef20dd24 --- /dev/null +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -0,0 +1,168 @@ +import logging +import sys +from typing import Any, Dict, List, Optional + +from hamilton import base +from hamilton.async_driver import AsyncDriver +from haystack import Document, component +from haystack_integrations.document_stores.qdrant import QdrantDocumentStore +from langfuse.decorators import observe + +from src.core.pipeline import BasicPipeline +from src.core.provider import DocumentStoreProvider, EmbedderProvider +from src.pipelines.common import ScoreFilter + +logger = logging.getLogger("wren-ai-service") + + +@component +class OutputFormatter: + @component.output_types( + documents=List[Optional[Dict]], + ) + def run(self, documents: List[Document]): + list = [] + + for doc in documents: + formatted = { + "instruction": doc.content, + "question": doc.meta.get("question", ""), + "instruction_id": doc.meta.get("instruction_id", ""), + } + list.append(formatted) + + return {"documents": list} + + +## Start of Pipeline +@observe(capture_input=False) +async def count_documents( + store: QdrantDocumentStore, project_id: Optional[str] = None +) -> int: + filters = ( + { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": project_id}, + ], + } + if project_id + else None + ) + document_count = await store.count_documents(filters=filters) + return document_count + + +@observe(capture_input=False, capture_output=False) +async def embedding(count_documents: int, query: str, embedder: Any) -> dict: + if count_documents: + return await embedder.run(query) + + return {} + + +@observe(capture_input=False) +async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: + if not embedding: + return {} + + filters = ( + { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": project_id}, + ], + } + if project_id + else None + ) + + res = await retriever.run( + query_embedding=embedding.get("embedding"), + filters=filters, + ) + return dict(documents=res.get("documents")) + + +@observe(capture_input=False) +def filtered_documents( + retrieval: dict, + score_filter: ScoreFilter, + similarity_threshold: float, + top_k: int, +) -> dict: + if not retrieval: + return {} + + return score_filter.run( + documents=retrieval.get("documents"), + score=similarity_threshold, + max_size=top_k, + ) + + +@observe(capture_input=False) +def formatted_output( + filtered_documents: dict, output_formatter: OutputFormatter +) -> dict: + if not filtered_documents: + return {"documents": []} + + documents = filtered_documents.get("documents") + return output_formatter.run(documents=documents) + + +## End of Pipeline + + +class Instructions(BasicPipeline): + def __init__( + self, + embedder_provider: EmbedderProvider, + document_store_provider: DocumentStoreProvider, + similarity_threshold: Optional[float] = 0.7, + top_k: Optional[int] = 10, + **kwargs, + ) -> None: + store = document_store_provider.get_store(dataset_name="instructions") + self._components = { + "store": store, + "embedder": embedder_provider.get_text_embedder(), + "retriever": document_store_provider.get_retriever( + document_store=store, + ), + "score_filter": ScoreFilter(), + "output_formatter": OutputFormatter(), + } + self._configs = { + "similarity_threshold": similarity_threshold, + "top_k": top_k, + } + + super().__init__( + AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult()) + ) + + @observe(name="Instructions Retrieval") + async def run(self, query: str, project_id: Optional[str] = None): + logger.info("Instructions Retrieval pipeline is running...") + return await self._pipe.execute( + ["formatted_output"], + inputs={ + "query": query, + "project_id": project_id or "", + **self._components, + **self._configs, + }, + ) + + +if __name__ == "__main__": + from src.pipelines.common import dry_run_pipeline + + dry_run_pipeline( + Instructions, + "instructions_retrieval", + query="how many customers live in EU?", + project_id="1", + ) From 9ee85042c58fff8d27b8b81a347282e3479fe889 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 15:49:28 +0800 Subject: [PATCH 05/30] chore: update config.yaml of docker and k8s for instructions indexing and retrieval pipe --- deployment/kustomizations/base/cm.yaml | 8 +++++++- docker/config.example.yaml | 6 ++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/deployment/kustomizations/base/cm.yaml b/deployment/kustomizations/base/cm.yaml index 32bc6854da..1602208ace 100644 --- a/deployment/kustomizations/base/cm.yaml +++ b/deployment/kustomizations/base/cm.yaml @@ -177,6 +177,12 @@ data: - name: sql_regeneration llm: litellm_llm.default engine: wren_ui + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.default + document_store: qdrant --- settings: @@ -192,4 +198,4 @@ data: langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: DEBUG - development: false \ No newline at end of file + development: false diff --git a/docker/config.example.yaml b/docker/config.example.yaml index 9e42f689f6..1c1a1c24f5 100644 --- a/docker/config.example.yaml +++ b/docker/config.example.yaml @@ -129,6 +129,12 @@ pipes: - name: sql_regeneration llm: litellm_llm.default engine: wren_ui + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.default + document_store: qdrant --- settings: From 5bbae3169bb35bf2b5e0df50c570f4a3bbc69b4e Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 16:36:17 +0800 Subject: [PATCH 06/30] feat: implement instructions service --- .../src/web/v1/services/instructions.py | 142 ++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 wren-ai-service/src/web/v1/services/instructions.py diff --git a/wren-ai-service/src/web/v1/services/instructions.py b/wren-ai-service/src/web/v1/services/instructions.py new file mode 100644 index 0000000000..d7b912716a --- /dev/null +++ b/wren-ai-service/src/web/v1/services/instructions.py @@ -0,0 +1,142 @@ +import logging +from typing import Dict, List, Literal, Optional + +from cachetools import TTLCache +from langfuse.decorators import observe +from pydantic import BaseModel + +from src.core.pipeline import BasicPipeline +from src.pipelines.indexing.instructions import Instruction +from src.utils import trace_metadata +from src.web.v1.services import MetadataTraceable + +logger = logging.getLogger("wren-ai-service") + + +class InstructionsService: + class Error(BaseModel): + code: Literal["OTHERS"] + message: str + + class Event(BaseModel, MetadataTraceable): + id: str + status: Literal["indexing", "deleting", "finished", "failed"] = "indexing" + error: Optional["InstructionsService.Error"] = None + trace_id: Optional[str] = None + + def __init__( + self, + pipelines: Dict[str, BasicPipeline], + maxsize: int = 1_000_000, + ttl: int = 120, + ): + self._pipelines = pipelines + self._cache: Dict[str, self.Event] = TTLCache(maxsize=maxsize, ttl=ttl) + + # todo: move it to utils for super class? + def _handle_exception( + self, + id: str, + error_message: str, + code: str = "OTHERS", + trace_id: Optional[str] = None, + ): + self._cache[id] = self.Event( + id=id, + status="failed", + error=self.Event.Error(code=code, message=error_message), + trace_id=trace_id, + ) + logger.error(error_message) + + class IndexRequest(BaseModel): + id: str + instructions: List[Instruction] + project_id: Optional[str] = None + + @observe(name="Index Instructions") + @trace_metadata + async def index( + self, + request: IndexRequest, + **kwargs, + ): + logger.info( + f"Request {request.id}: Instructions Indexing process is running..." + ) + trace_id = kwargs.get("trace_id") + + try: + await self._pipelines["instructions_indexing"].run( + project_id=request.project_id, + instructions=request.instructions, + ) + + self._cache[request.id] = self.Event( + id=request.id, + status="finished", + trace_id=trace_id, + ) + + except Exception as e: + self._handle_exception( + request.id, + f"An error occurred during instructions indexing: {str(e)}", + trace_id=trace_id, + ) + + return self._cache[request.id].with_metadata() + + class DeleteRequest(BaseModel): + id: str + instruction_ids: List[str] + project_id: Optional[str] = None + + @observe(name="Delete Instructions") + @trace_metadata + async def delete( + self, + request: DeleteRequest, + **kwargs, + ): + logger.info( + f"Request {request.id}: Instructions Deletion process is running..." + ) + trace_id = kwargs.get("trace_id") + + try: + instructions = [Instruction(id=id) for id in request.instruction_ids] + await self._pipelines["instructions_indexing"].clean( + instructions=instructions, project_id=request.project_id + ) + + self._cache[request.id] = self.Event( + id=request.id, + status="finished", + trace_id=trace_id, + ) + except Exception as e: + self._handle_exception( + request.id, + f"Failed to delete instructions: {e}", + trace_id=trace_id, + ) + + return self._cache[request.id].with_metadata() + + def __getitem__(self, id: str) -> Event: + response = self._cache.get(id) + + if response is None: + message = f"Instructions Event with ID '{id}' not found." + logger.exception(message) + return self.Event( + id=id, + status="failed", + error=self.Event.Error(code="OTHERS", message=message), + ) + + return response + + def __setitem__(self, id: str, value: Event): + self._cache[id] = value From acde02883483ace91f6982a07de674904d60acf5 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 17:32:33 +0800 Subject: [PATCH 07/30] feat: add instruction service into global --- wren-ai-service/src/config.py | 2 ++ wren-ai-service/src/globals.py | 14 ++++++++++++++ wren-ai-service/src/pipelines/indexing/__init__.py | 2 ++ .../src/pipelines/retrieval/__init__.py | 2 ++ wren-ai-service/src/web/v1/services/__init__.py | 2 ++ 5 files changed, 22 insertions(+) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index 21ab4d71ba..9c0ae408bf 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -33,6 +33,8 @@ class Settings(BaseSettings): 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) + instructions_similarity_threshold: float = Field(default=0.7) + instructions_top_k: int = Field(default=10) # generation config allow_intent_classification: bool = Field(default=True) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index cc64e79890..6b13df43cb 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -26,6 +26,7 @@ class ServiceContainer: sql_expansion_service: services.SqlExpansionService sql_pairs_service: services.SqlPairsService sql_question_service: services.SqlQuestionService + instructions_service: services.InstructionsService @dataclass @@ -93,6 +94,11 @@ def create_service_container( 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, @@ -237,6 +243,14 @@ def create_service_container( }, **query_cache, ), + instructions_service=services.InstructionsService( + pipelines={ + "instructions_indexing": indexing.Instructions( + **pipe_components["instructions_indexing"], + ) + }, + **query_cache, + ), ) diff --git a/wren-ai-service/src/pipelines/indexing/__init__.py b/wren-ai-service/src/pipelines/indexing/__init__.py index 71ef286c53..885c568931 100644 --- a/wren-ai-service/src/pipelines/indexing/__init__.py +++ b/wren-ai-service/src/pipelines/indexing/__init__.py @@ -90,6 +90,7 @@ async def run( # Put the pipelines imports here to avoid circular imports and make them accessible directly to the rest of packages from .db_schema import DBSchema # noqa: E402 from .historical_question import HistoricalQuestion # noqa: E402 +from .instructions import Instructions # noqa: E402 from .sql_pairs import SqlPairs # noqa: E402 from .table_description import TableDescription # noqa: E402 @@ -98,4 +99,5 @@ async def run( "TableDescription", "HistoricalQuestion", "SqlPairs", + "Instructions", ] diff --git a/wren-ai-service/src/pipelines/retrieval/__init__.py b/wren-ai-service/src/pipelines/retrieval/__init__.py index 82e319b4c6..5613b0a694 100644 --- a/wren-ai-service/src/pipelines/retrieval/__init__.py +++ b/wren-ai-service/src/pipelines/retrieval/__init__.py @@ -1,4 +1,5 @@ from .historical_question_retrieval import HistoricalQuestionRetrieval +from .instructions import Instructions from .preprocess_sql_data import PreprocessSqlData from .retrieval import Retrieval from .sql_executor import SQLExecutor @@ -10,4 +11,5 @@ "Retrieval", "SQLExecutor", "SqlPairsRetrieval", + "Instructions", ] diff --git a/wren-ai-service/src/web/v1/services/__init__.py b/wren-ai-service/src/web/v1/services/__init__.py index 3a9839d3cb..55c1ed9ec9 100644 --- a/wren-ai-service/src/web/v1/services/__init__.py +++ b/wren-ai-service/src/web/v1/services/__init__.py @@ -63,6 +63,7 @@ def serialize(self): from .ask_details import AskDetailsService # noqa: E402 from .chart import ChartService # noqa: E402 from .chart_adjustment import ChartAdjustmentService # noqa: E402 +from .instructions import InstructionsService # noqa: E402 from .question_recommendation import QuestionRecommendation # noqa: E402 from .relationship_recommendation import RelationshipRecommendation # noqa: E402 from .semantics_description import SemanticsDescription # noqa: E402 @@ -85,4 +86,5 @@ def serialize(self): "SqlExpansionService", "SqlPairsService", "SqlQuestionService", + "InstructionsService", ] From 46152f2a1bd1c22d03e03e531ec18725017d9181 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 17:47:30 +0800 Subject: [PATCH 08/30] feat: implement instruction endpoint --- .../src/web/v1/routers/__init__.py | 3 +- .../src/web/v1/routers/instructions.py | 155 ++++++++++++++++++ .../src/web/v1/services/instructions.py | 40 ++--- 3 files changed, 177 insertions(+), 21 deletions(-) create mode 100644 wren-ai-service/src/web/v1/routers/instructions.py diff --git a/wren-ai-service/src/web/v1/routers/__init__.py b/wren-ai-service/src/web/v1/routers/__init__.py index dc108201cf..d67687a3a1 100644 --- a/wren-ai-service/src/web/v1/routers/__init__.py +++ b/wren-ai-service/src/web/v1/routers/__init__.py @@ -5,6 +5,7 @@ ask_details, chart, chart_adjustment, + instructions, question_recommendation, relationship_recommendation, semantics_description, @@ -28,4 +29,4 @@ router.include_router(chart_adjustment.router) router.include_router(sql_pairs.router) router.include_router(sql_question.router) -# connected subrouter +router.include_router(instructions.router) diff --git a/wren-ai-service/src/web/v1/routers/instructions.py b/wren-ai-service/src/web/v1/routers/instructions.py new file mode 100644 index 0000000000..d2a8468cd5 --- /dev/null +++ b/wren-ai-service/src/web/v1/routers/instructions.py @@ -0,0 +1,155 @@ +import uuid +from dataclasses import asdict +from typing import List, Literal, Optional + +from fastapi import APIRouter, BackgroundTasks, Depends, Response +from pydantic import BaseModel + +from src.globals import ( + ServiceContainer, + ServiceMetadata, + get_service_container, + get_service_metadata, +) +from src.pipelines.indexing.instructions import Instruction +from src.web.v1.services import InstructionsService + +router = APIRouter() + + +""" +Instructions Router + +This router handles endpoints related to preparing and managing instructions. + +Endpoints: +1. POST /instructions + - Indexes instructions for retrieval + - Request body: PostRequest + { + "instructions": [ # List of instructions + { + "id": "unique-id", # Unique identifier for the instruction + "content": "Instruction content", # Content of the instruction + "metadata": {} # Optional metadata for the instruction + } + ], + "project_id": "project-id" # Optional project ID + } + - Response: PostResponse + { + "event_id": "unique-uuid" # Unique identifier for tracking indexing + } + +2. DELETE /instructions + - Deletes specified instructions by their IDs + - Request body: DeleteRequest + { + "instruction_ids": ["id1", "id2"], # List of instruction IDs to delete + "project_id": "project-id" # Optional project ID + } + - Response: None or Error object if failed + +3. GET /instructions/{event_id} + - Retrieves status of instructions indexing/deletion + - Path parameter: event_id (str) + - Response: GetResponse + { + "event_id": "unique-uuid", # Unique identifier + "status": "indexing" | "deleting" | "finished" | "failed", + "error": { # Present only if status is "failed" + "code": "OTHERS", + "message": "Error description" + }, + "trace_id": "trace-id" # Optional trace ID for tracking + } + +The instructions indexing and deletion are asynchronous processes. The POST endpoint +initiates the operation and returns immediately with an ID. The GET endpoint can then be used +to check the status and result. + +Usage: +1. Send a POST/DELETE request to start the operation +2. Use the returned ID to poll the GET endpoint until status is "finished" or "failed" + +Note: The actual processing is performed in the background using FastAPI's BackgroundTasks. +Results are cached with a TTL defined in the service configuration. +""" + + +class PostRequest(BaseModel): + instructions: List[Instruction] + project_id: Optional[str] = None + + +class PostResponse(BaseModel): + event_id: str + + +@router.post("/instructions") +async def index( + request: PostRequest, + background_tasks: BackgroundTasks, + service_container: ServiceContainer = Depends(get_service_container), + service_metadata: ServiceMetadata = Depends(get_service_metadata), +) -> PostResponse: + event_id = str(uuid.uuid4()) + service = service_container.instructions_service + service[event_id] = InstructionsService.Event(event_id=event_id) + + index_request = InstructionsService.IndexRequest( + event_id=event_id, **request.model_dump() + ) + + background_tasks.add_task( + service.index, + index_request, + service_metadata=asdict(service_metadata), + ) + return PostResponse(event_id=event_id) + + +class DeleteRequest(BaseModel): + instruction_ids: List[str] + project_id: Optional[str] = None + + +@router.delete("/instructions") +async def delete( + request: DeleteRequest, + response: Response, + service_container: ServiceContainer = Depends(get_service_container), + service_metadata: ServiceMetadata = Depends(get_service_metadata), +) -> None | InstructionsService.Error: + event_id = str(uuid.uuid4()) + service = service_container.instructions_service + service[event_id] = InstructionsService.Event(event_id=event_id, status="deleting") + + delete_request = InstructionsService.DeleteRequest( + event_id=event_id, + **request.model_dump(), + ) + + await service.delete(delete_request, service_metadata=asdict(service_metadata)) + + event: InstructionsService.Event = service[event_id] + + if event.status == "failed": + response.status_code = 500 + return event.error + + +class GetResponse(BaseModel): + event_id: str + status: Literal["indexing", "deleting", "finished", "failed"] + error: Optional[dict] + trace_id: Optional[str] + + +@router.get("/instructions/{event_id}") +async def get( + event_id: str, + container: ServiceContainer = Depends(get_service_container), +) -> GetResponse: + event: InstructionsService.Event = container.instructions_service[event_id] + return GetResponse(**event.model_dump()) diff --git a/wren-ai-service/src/web/v1/services/instructions.py b/wren-ai-service/src/web/v1/services/instructions.py index d7b912716a..b3b935e956 100644 --- a/wren-ai-service/src/web/v1/services/instructions.py +++ b/wren-ai-service/src/web/v1/services/instructions.py @@ -19,7 +19,7 @@ class Error(BaseModel): message: str class Event(BaseModel, MetadataTraceable): - id: str + event_id: str status: Literal["indexing", "deleting", "finished", "failed"] = "indexing" error: Optional["InstructionsService.Error"] = None trace_id: Optional[str] = None @@ -42,7 +42,7 @@ def _handle_exception( trace_id: Optional[str] = None, ): self._cache[id] = self.Event( - id=id, + event_id=id, status="failed", error=self.Event.Error(code=code, message=error_message), trace_id=trace_id, @@ -50,7 +50,7 @@ def _handle_exception( logger.error(error_message) class IndexRequest(BaseModel): - id: str + event_id: str instructions: List[Instruction] project_id: Optional[str] = None @@ -62,7 +62,7 @@ async def index( **kwargs, ): logger.info( - f"Request {request.id}: Instructions Indexing process is running..." + f"Request {request.event_id}: Instructions Indexing process is running..." ) trace_id = kwargs.get("trace_id") @@ -72,23 +72,23 @@ async def index( instructions=request.instructions, ) - self._cache[request.id] = self.Event( - id=request.id, + self._cache[request.event_id] = self.Event( + event_id=request.event_id, status="finished", trace_id=trace_id, ) except Exception as e: self._handle_exception( - request.id, + request.event_id, f"An error occurred during instructions indexing: {str(e)}", trace_id=trace_id, ) - return self._cache[request.id].with_metadata() + return self._cache[request.event_id].with_metadata() class DeleteRequest(BaseModel): - id: str + event_id: str instruction_ids: List[str] project_id: Optional[str] = None @@ -100,7 +100,7 @@ async def delete( **kwargs, ): logger.info( - f"Request {request.id}: Instructions Deletion process is running..." + f"Request {request.event_id}: Instructions Deletion process is running..." ) trace_id = kwargs.get("trace_id") @@ -110,33 +110,33 @@ async def delete( instructions=instructions, project_id=request.project_id ) - self._cache[request.id] = self.Event( - id=request.id, + self._cache[request.event_id] = self.Event( + event_id=request.event_id, status="finished", trace_id=trace_id, ) except Exception as e: self._handle_exception( - request.id, + request.event_id, f"Failed to delete instructions: {e}", trace_id=trace_id, ) - return self._cache[request.id].with_metadata() + return self._cache[request.event_id].with_metadata() - def __getitem__(self, id: str) -> Event: - response = self._cache.get(id) + def __getitem__(self, event_id: str) -> Event: + response = self._cache.get(event_id) if response is None: - message = f"Instructions Event with ID '{id}' not found." + message = f"Instructions Event with ID '{event_id}' not found." logger.exception(message) return self.Event( - id=id, + event_id=event_id, status="failed", error=self.Event.Error(code="OTHERS", message=message), ) return response - def __setitem__(self, id: str, value: Event): - self._cache[id] = value + def __setitem__(self, event_id: str, value: Event): + self._cache[event_id] = value From dd472d88616210b95ad471e9243fcfcee37cf63b Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 17:53:13 +0800 Subject: [PATCH 09/30] feat: trigger delete instruction from vector store when delete semantics --- wren-ai-service/src/globals.py | 3 +++ wren-ai-service/src/pipelines/indexing/instructions.py | 2 +- .../src/web/v1/services/semantics_preparation.py | 7 ++++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index 6b13df43cb..48861f3613 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -68,6 +68,9 @@ def create_service_container( **pipe_components["sql_pairs_indexing"], sql_pairs_path=settings.sql_pairs_path, ), + "instructions": indexing.Instructions( + **pipe_components["instructions_indexing"], + ), }, **query_cache, ), diff --git a/wren-ai-service/src/pipelines/indexing/instructions.py b/wren-ai-service/src/pipelines/indexing/instructions.py index 2a8581777d..6cfd8418e2 100644 --- a/wren-ai-service/src/pipelines/indexing/instructions.py +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -162,7 +162,7 @@ async def run( @observe(name="Clean Documents for Instructions") async def clean( self, - instructions: List[Instruction], + instructions: List[Instruction] = [], project_id: Optional[str] = None, delete_all: bool = False, ) -> None: diff --git a/wren-ai-service/src/web/v1/services/semantics_preparation.py b/wren-ai-service/src/web/v1/services/semantics_preparation.py index ffd4449363..a6cf62a99f 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -135,7 +135,7 @@ def get_prepare_semantics_status( @observe(name="Delete Semantics Documents") @trace_metadata - async def delete_semantics(self, project_id: str): + async def delete_semantics(self, project_id: str, **kwargs): logger.info(f"Project ID: {project_id}, Deleting semantics documents...") tasks = [ @@ -147,6 +147,11 @@ async def delete_semantics(self, project_id: str): project_id=project_id, delete_all=True, ) + ] + [ + self._pipelines["instructions"].clean( + project_id=project_id, + delete_all=True, + ) ] await asyncio.gather(*tasks) From bacd07ff3f3aaec85160362013ac50f68beeafb8 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 18:00:25 +0800 Subject: [PATCH 10/30] refactor: simplify the code base for deleting semantics --- .../src/pipelines/indexing/instructions.py | 3 +- .../src/pipelines/indexing/sql_pairs.py | 5 ++- .../web/v1/services/semantics_preparation.py | 33 ++++++++----------- 3 files changed, 17 insertions(+), 24 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/instructions.py b/wren-ai-service/src/pipelines/indexing/instructions.py index 6cfd8418e2..59c33e2389 100644 --- a/wren-ai-service/src/pipelines/indexing/instructions.py +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -96,7 +96,7 @@ async def embedding( async def clean( cleaner: InstructionsCleaner, instructions: List[Instruction], - embedding: Dict[str, Any], + embedding: Dict[str, Any] = {}, project_id: Optional[str] = "", delete_all: bool = False, ) -> Dict[str, Any]: @@ -168,7 +168,6 @@ async def clean( ) -> None: await clean( instructions=instructions, - embedding={"documents": []}, cleaner=self._components["cleaner"], project_id=project_id, delete_all=delete_all, diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index a5e8e68321..53f0aa5546 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -125,7 +125,7 @@ async def embedding( async def clean( cleaner: SqlPairsCleaner, sql_pairs: List[SqlPair], - embedding: Dict[str, Any], + embedding: Dict[str, Any] = {}, project_id: Optional[str] = "", delete_all: bool = False, ) -> Dict[str, Any]: @@ -215,13 +215,12 @@ async def run( @observe(name="Clean Documents for SQL Pairs") async def clean( self, - sql_pairs: List[SqlPair], + sql_pairs: List[SqlPair] = [], project_id: Optional[str] = None, delete_all: bool = False, ) -> None: await clean( sql_pairs=sql_pairs, - embedding={"documents": []}, cleaner=self._components["cleaner"], project_id=project_id, delete_all=delete_all, diff --git a/wren-ai-service/src/web/v1/services/semantics_preparation.py b/wren-ai-service/src/web/v1/services/semantics_preparation.py index a6cf62a99f..661fda5121 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -89,22 +89,22 @@ async def prepare_semantics( await asyncio.gather(*tasks) - self._prepare_semantics_statuses[ - prepare_semantics_request.mdl_hash - ] = SemanticsPreparationStatusResponse( - status="finished", + self._prepare_semantics_statuses[prepare_semantics_request.mdl_hash] = ( + SemanticsPreparationStatusResponse( + status="finished", + ) ) except Exception as e: logger.exception(f"Failed to prepare semantics: {e}") - self._prepare_semantics_statuses[ - prepare_semantics_request.mdl_hash - ] = SemanticsPreparationStatusResponse( - status="failed", - error=SemanticsPreparationStatusResponse.SemanticsPreparationError( - code="OTHERS", - message=f"Failed to prepare semantics: {e}", - ), + self._prepare_semantics_statuses[prepare_semantics_request.mdl_hash] = ( + SemanticsPreparationStatusResponse( + status="failed", + error=SemanticsPreparationStatusResponse.SemanticsPreparationError( + code="OTHERS", + message=f"Failed to prepare semantics: {e}", + ), + ) ) results["metadata"]["error_type"] = "INDEXING_FAILED" @@ -142,16 +142,11 @@ async def delete_semantics(self, project_id: str, **kwargs): self._pipelines[name].clean(project_id=project_id) for name in ["db_schema", "historical_question", "table_description"] ] + [ - self._pipelines["sql_pairs"].clean( - sql_pairs=[], - project_id=project_id, - delete_all=True, - ) - ] + [ - self._pipelines["instructions"].clean( + self._pipelines[name].clean( project_id=project_id, delete_all=True, ) + for name in ["sql_pairs", "instructions"] ] await asyncio.gather(*tasks) From 13c90160e5cf676713ffdf3fb4a5f37c2567f0ad Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 6 Mar 2025 19:23:27 +0800 Subject: [PATCH 11/30] feat: retrieve the default is true instructions --- .../src/pipelines/indexing/instructions.py | 2 + .../src/pipelines/retrieval/instructions.py | 54 +++++++++++++------ 2 files changed, 40 insertions(+), 16 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/instructions.py b/wren-ai-service/src/pipelines/indexing/instructions.py index 59c33e2389..298c6baf75 100644 --- a/wren-ai-service/src/pipelines/indexing/instructions.py +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -21,6 +21,7 @@ class Instruction(BaseModel): id: str instruction: str = "" question: str = "" + is_default: bool = False @component @@ -39,6 +40,7 @@ def run(self, instructions: list[Instruction], project_id: Optional[str] = ""): meta={ "instruction_id": instruction.id, "question": instruction.question, + "is_default": instruction.is_default, **addition, }, content=instruction.instruction, diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 5bef20dd24..054bad164e 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -66,16 +66,17 @@ async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: if not embedding: return {} - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": False}, + ], + } + + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) res = await retriever.run( query_embedding=embedding.get("embedding"), @@ -101,15 +102,36 @@ def filtered_documents( ) +@observe(capture_input=False) +def default_instructions(store: QdrantDocumentStore, project_id: str) -> list[Document]: + filters = { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": True}, + ], + } + + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + + return store.filter_documents(filters=filters) + + @observe(capture_input=False) def formatted_output( - filtered_documents: dict, output_formatter: OutputFormatter + default_instructions: list[Document], + filtered_documents: dict, + output_formatter: OutputFormatter, ) -> dict: - if not filtered_documents: + if not filtered_documents and not default_instructions: return {"documents": []} - documents = filtered_documents.get("documents") - return output_formatter.run(documents=documents) + # todo: default instruction also need to be formatted + documents = output_formatter.run(documents=filtered_documents.get("documents")) + documents["documents"].extend(default_instructions) + return documents ## End of Pipeline @@ -163,6 +185,6 @@ async def run(self, query: str, project_id: Optional[str] = None): dry_run_pipeline( Instructions, "instructions_retrieval", - query="how many customers live in EU?", - project_id="1", + query="hello", + project_id="string", ) From a21040bd36d4ad01e6028bf52146e2158e72245e Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Fri, 7 Mar 2025 11:35:51 +0800 Subject: [PATCH 12/30] chore: add a comment to explain default instruction and correct the document content for instruction --- wren-ai-service/src/pipelines/indexing/instructions.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/instructions.py b/wren-ai-service/src/pipelines/indexing/instructions.py index 298c6baf75..0fdf36951f 100644 --- a/wren-ai-service/src/pipelines/indexing/instructions.py +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -21,6 +21,7 @@ class Instruction(BaseModel): id: str instruction: str = "" question: str = "" + # This is used to identify the default instruction needed to be retrieved for the project is_default: bool = False @@ -32,18 +33,17 @@ def run(self, instructions: list[Instruction], project_id: Optional[str] = ""): addition = {"project_id": project_id} if project_id else {} - # todo: for the prototype, the content using instruction directly return { "documents": [ Document( id=str(uuid.uuid4()), meta={ "instruction_id": instruction.id, - "question": instruction.question, + "instruction": instruction.instruction, "is_default": instruction.is_default, **addition, }, - content=instruction.instruction, + content=instruction.question, ) for instruction in instructions ] From 1388eae226fbfd1332319e20d9032fa1b1382896 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Fri, 7 Mar 2025 11:54:33 +0800 Subject: [PATCH 13/30] feat: flatten the request questions to one-on-one question instruction pair --- .../src/web/v1/routers/instructions.py | 3 +-- .../src/web/v1/services/instructions.py | 22 +++++++++++++++++-- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/web/v1/routers/instructions.py b/wren-ai-service/src/web/v1/routers/instructions.py index d2a8468cd5..d810c582cd 100644 --- a/wren-ai-service/src/web/v1/routers/instructions.py +++ b/wren-ai-service/src/web/v1/routers/instructions.py @@ -11,7 +11,6 @@ get_service_container, get_service_metadata, ) -from src.pipelines.indexing.instructions import Instruction from src.web.v1.services import InstructionsService router = APIRouter() @@ -78,7 +77,7 @@ class PostRequest(BaseModel): - instructions: List[Instruction] + instructions: List[InstructionsService.Instruction] project_id: Optional[str] = None diff --git a/wren-ai-service/src/web/v1/services/instructions.py b/wren-ai-service/src/web/v1/services/instructions.py index b3b935e956..27e9e1eb45 100644 --- a/wren-ai-service/src/web/v1/services/instructions.py +++ b/wren-ai-service/src/web/v1/services/instructions.py @@ -14,6 +14,13 @@ class InstructionsService: + class Instruction(BaseModel): + id: str + instruction: str + questions: List[str] + # This is used to identify the default instruction needed to be retrieved for the project + is_default: bool = False + class Error(BaseModel): code: Literal["OTHERS"] message: str @@ -51,7 +58,7 @@ def _handle_exception( class IndexRequest(BaseModel): event_id: str - instructions: List[Instruction] + instructions: List["InstructionsService.Instruction"] project_id: Optional[str] = None @observe(name="Index Instructions") @@ -67,9 +74,20 @@ async def index( trace_id = kwargs.get("trace_id") try: + instructions = [ + Instruction( + id=instruction.id, + instruction=instruction.instruction, + question=question, + is_default=instruction.is_default, + ) + for instruction in request.instructions + for question in instruction.questions + ] + await self._pipelines["instructions_indexing"].run( project_id=request.project_id, - instructions=request.instructions, + instructions=instructions, ) self._cache[request.event_id] = self.Event( From 0f8cfb9eb768f1564c80a6067f213a8ad2effa73 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Fri, 7 Mar 2025 16:35:02 +0800 Subject: [PATCH 14/30] feat: retrieve the instructions for ask process --- .../generation/followup_sql_generation.py | 4 ++++ .../src/pipelines/generation/sql_generation.py | 4 ++++ .../src/pipelines/generation/utils/sql.py | 18 ++++++++++++------ .../src/pipelines/retrieval/instructions.py | 1 - wren-ai-service/src/web/v1/services/ask.py | 11 +++++++++++ 5 files changed, 31 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py index a89602653b..d88c29d4fa 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -75,6 +75,7 @@ def prompt( configuration: Configuration, prompt_builder: PromptBuilder, sql_samples: List[Dict] | None = None, + instructions: List[Dict] | None = None, has_calculated_field: bool = False, has_metric: bool = False, ) -> dict: @@ -91,6 +92,7 @@ def prompt( has_calculated_field, has_metric, sql_samples, + instructions, ), current_time=configuration.show_current_time(), sql_samples=sql_samples, @@ -155,6 +157,7 @@ async def run( histories: list[AskHistory], configuration: Configuration = Configuration(), sql_samples: List[Dict] | None = None, + instructions: List[Dict] | None = None, project_id: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, @@ -170,6 +173,7 @@ async def run( "project_id": project_id, "configuration": configuration, "sql_samples": sql_samples, + "instructions": instructions, "has_calculated_field": has_calculated_field, "has_metric": has_metric, **self._components, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 8ca78fa39f..bca59ff0f9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -65,6 +65,7 @@ def prompt( sql_generation_reasoning: str | None = None, configuration: Configuration | None = None, sql_samples: List[Dict] | None = None, + instructions: List[Dict] | None = None, has_calculated_field: bool = False, has_metric: bool = False, ) -> dict: @@ -77,6 +78,7 @@ def prompt( has_calculated_field, has_metric, sql_samples, + instructions, ), sql_samples=sql_samples, current_time=configuration.show_current_time(), @@ -143,6 +145,7 @@ async def run( sql_generation_reasoning: str | None = None, configuration: Configuration = Configuration(), sql_samples: List[Dict] | None = None, + instructions: List[Dict] | None = None, project_id: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, @@ -155,6 +158,7 @@ async def run( "documents": contexts, "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, + "instructions": instructions, "project_id": project_id, "configuration": configuration, "has_calculated_field": has_calculated_field, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 09e0e9ad9b..fd5d44ea74 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -534,22 +534,28 @@ def construct_instructions( has_calculated_field: bool = False, has_metric: bool = False, sql_samples: list | None = None, + instructions: list | None = None, ): if sql_samples is None: sql_samples = [] - instructions = "" + _instructions = "" if configuration: if configuration.fiscal_year: - instructions += f"\n- For calendar year related computation, it should be started from {configuration.fiscal_year.start} to {configuration.fiscal_year.end}\n\n" + _instructions += f"\n- For calendar year related computation, it should be started from {configuration.fiscal_year.start} to {configuration.fiscal_year.end}\n\n" if has_calculated_field: - instructions += calculated_field_instructions + _instructions += calculated_field_instructions if has_metric: - instructions += metric_instructions + _instructions += metric_instructions if sql_samples: - instructions += sql_samples_instructions + _instructions += sql_samples_instructions + if instructions: + # todo: refactor the format of the instructions + _instructions += "\n\n".join( + [f"# {instruction.get('instruction')}\n\n" for instruction in instructions] + ) - return instructions + return _instructions class SqlGenerationResult(BaseModel): diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 054bad164e..e21a1f10fc 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -128,7 +128,6 @@ def formatted_output( if not filtered_documents and not default_instructions: return {"documents": []} - # todo: default instruction also need to be formatted documents = output_formatter.run(documents=filtered_documents.get("documents")) documents["documents"].extend(default_instructions) return documents diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e5b6238ae0..159a066b16 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -366,6 +366,15 @@ async def ask( ) )["formatted_output"].get("documents", []) + # todo: consider to retireve at the same time with sql_samples + # todo: should we also put the instructioons into reasoning pipeline? + instructions = ( + await self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ) + )["formatted_output"].get("documents", []) + sql_generation_reasoning = ( await self._pipelines["sql_generation_reasoning"].run( query=user_query, @@ -416,6 +425,7 @@ async def ask( project_id=ask_request.project_id, configuration=ask_request.configurations, sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, ) @@ -429,6 +439,7 @@ async def ask( project_id=ask_request.project_id, configuration=ask_request.configurations, sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, ) From c31477e329dac91a527a8f9d120bf58027b958a9 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Fri, 7 Mar 2025 16:45:38 +0800 Subject: [PATCH 15/30] feat: correct the instruction format for prompt and fix the wrong pass for object --- wren-ai-service/src/pipelines/generation/utils/sql.py | 3 ++- wren-ai-service/src/pipelines/retrieval/instructions.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index fd5d44ea74..841b166504 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -551,8 +551,9 @@ def construct_instructions( _instructions += sql_samples_instructions if instructions: # todo: refactor the format of the instructions + _instructions += "\n\n#### Instructions for SQL Generation ####\n\n" _instructions += "\n\n".join( - [f"# {instruction.get('instruction')}\n\n" for instruction in instructions] + [f"{instruction.get('instruction')}\n\n" for instruction in instructions] ) return _instructions diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index e21a1f10fc..2194d5da6a 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -25,8 +25,8 @@ def run(self, documents: List[Document]): for doc in documents: formatted = { - "instruction": doc.content, - "question": doc.meta.get("question", ""), + "instruction": doc.meta.get("instruction", ""), + "question": doc.content, "instruction_id": doc.meta.get("instruction_id", ""), } list.append(formatted) From 37b50027a2683b4150e32a7bb52196a09e437eed Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Fri, 7 Mar 2025 16:59:04 +0800 Subject: [PATCH 16/30] chore: set the pipeline config to default embedder --- wren-ai-service/docs/config_examples/config.azure.yaml | 2 +- wren-ai-service/docs/config_examples/config.deepseek.yaml | 2 +- .../docs/config_examples/config.google_ai_studio.yaml | 2 +- wren-ai-service/docs/config_examples/config.groq.yaml | 2 +- wren-ai-service/docs/config_examples/config.ollama.yaml | 2 +- wren-ai-service/tools/config/config.example.yaml | 2 +- wren-ai-service/tools/config/config.full.yaml | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/docs/config_examples/config.azure.yaml b/wren-ai-service/docs/config_examples/config.azure.yaml index a8f2da1186..7bd5241f79 100644 --- a/wren-ai-service/docs/config_examples/config.azure.yaml +++ b/wren-ai-service/docs/config_examples/config.azure.yaml @@ -136,7 +136,7 @@ pipes: embedder: litellm_embedder.default document_store: qdrant - name: instructions_retrieval - embedder: litellm_embedder.azure/text-embedding-ada-002 + embedder: litellm_embedder.default document_store: qdrant --- diff --git a/wren-ai-service/docs/config_examples/config.deepseek.yaml b/wren-ai-service/docs/config_examples/config.deepseek.yaml index 5b18bd2cc6..59b4d4be97 100644 --- a/wren-ai-service/docs/config_examples/config.deepseek.yaml +++ b/wren-ai-service/docs/config_examples/config.deepseek.yaml @@ -146,7 +146,7 @@ pipes: embedder: litellm_embedder.default document_store: qdrant - name: instructions_retrieval - embedder: litellm_embedder.text-embedding-3-large + embedder: litellm_embedder.default document_store: qdrant --- diff --git a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml index 680f07942e..a540e795ff 100644 --- a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml +++ b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml @@ -135,7 +135,7 @@ pipes: embedder: litellm_embedder.default document_store: qdrant - name: instructions_retrieval - embedder: litellm_embedder.gemini/text-embedding-004 + embedder: litellm_embedder.default document_store: qdrant --- diff --git a/wren-ai-service/docs/config_examples/config.groq.yaml b/wren-ai-service/docs/config_examples/config.groq.yaml index d4a8fe915f..47c20a1b7c 100644 --- a/wren-ai-service/docs/config_examples/config.groq.yaml +++ b/wren-ai-service/docs/config_examples/config.groq.yaml @@ -128,7 +128,7 @@ pipes: embedder: litellm_embedder.default document_store: qdrant - name: instructions_retrieval - embedder: litellm_embedder.text-embedding-3-large + embedder: litellm_embedder.default document_store: qdrant --- diff --git a/wren-ai-service/docs/config_examples/config.ollama.yaml b/wren-ai-service/docs/config_examples/config.ollama.yaml index 26c9faf53c..6d8055c91c 100644 --- a/wren-ai-service/docs/config_examples/config.ollama.yaml +++ b/wren-ai-service/docs/config_examples/config.ollama.yaml @@ -125,7 +125,7 @@ pipes: embedder: litellm_embedder.default document_store: qdrant - name: instructions_retrieval - embedder: litellm_embedder.openai/nomic-embed-text + embedder: litellm_embedder.default document_store: qdrant --- diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index e4b2a5914d..2e030cda33 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -149,7 +149,7 @@ pipes: embedder: litellm_embedder.default document_store: qdrant - name: instructions_retrieval - embedder: litellm_embedder.text-embedding-3-large + embedder: litellm_embedder.default document_store: qdrant --- diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 001f43ea34..b4bf7ece25 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -149,7 +149,7 @@ pipes: embedder: litellm_embedder.default document_store: qdrant - name: instructions_retrieval - embedder: litellm_embedder.text-embedding-3-large + embedder: litellm_embedder.default document_store: qdrant --- From ca47341befcdc2f316f4de158c4495ff4b3b2695 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Fri, 7 Mar 2025 17:02:19 +0800 Subject: [PATCH 17/30] chore: early return when document count is 0 --- wren-ai-service/src/pipelines/retrieval/instructions.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 2194d5da6a..925c20d426 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -103,7 +103,12 @@ def filtered_documents( @observe(capture_input=False) -def default_instructions(store: QdrantDocumentStore, project_id: str) -> list[Document]: +def default_instructions( + count_documents: int, store: QdrantDocumentStore, project_id: str +) -> list[Document]: + if not count_documents: + return [] + filters = { "operator": "AND", "conditions": [ From 0d2f5ec64e0ff7e3572d67e4d7d864261ce0d10e Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Fri, 7 Mar 2025 17:06:59 +0800 Subject: [PATCH 18/30] feat: add test case for instructions indexing --- .../pipelines/indexing/test_instructions.py | 140 ++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py new file mode 100644 index 0000000000..e242392b29 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py @@ -0,0 +1,140 @@ +import pytest + +from src.config import settings +from src.core.provider import DocumentStoreProvider +from src.pipelines.indexing import Instructions +from src.pipelines.indexing.instructions import Instruction +from src.providers import generate_components + + +@pytest.mark.asyncio +async def test_instructions_indexing(): + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + recreate_index=True, + ) + + instructions = Instructions(**pipe_components["instructions_indexing"]) + test_instructions = [ + Instruction( + id="test-id-1", + instruction="This is a test instruction", + question="What is the test question?", + is_default=False, + ), + Instruction( + id="test-id-2", + instruction="This is another test instruction", + question="What is another test question?", + is_default=True, + ), + ] + + await instructions.run( + project_id="fake-id", + instructions=test_instructions, + ) + + assert await store.count_documents() == 2 + documents = store.filter_documents() + for document in documents: + assert document.content, "content should not be empty" + assert document.meta, "meta should not be empty" + assert document.meta.get("instruction_id"), "instruction_id should be in meta" + assert document.meta.get("instruction"), "instruction should be in meta" + assert "is_default" in document.meta, "is_default should be in meta" + + +@pytest.mark.asyncio +async def test_instructions_indexing_with_multiple_project_ids(): + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + recreate_index=True, + ) + + instructions = Instructions(**pipe_components["instructions_indexing"]) + test_instructions = [ + Instruction( + id="test-id-1", + instruction="This is a test instruction", + question="What is the test question?", + ), + ] + + await instructions.run( + project_id="fake-id", + instructions=test_instructions, + ) + + await instructions.run( + project_id="fake-id-2", + instructions=test_instructions, + ) + + assert await store.count_documents() == 2 + documents = store.filter_documents( + filters={ + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "fake-id"}, + ], + } + ) + assert len(documents) == 1 + + +@pytest.mark.asyncio +async def test_instructions_deletion(): + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + recreate_index=True, + ) + + instructions = Instructions(**pipe_components["instructions_indexing"]) + test_instructions = [ + Instruction( + id="test-id-1", + instruction="This is a test instruction", + question="What is the test question?", + ), + Instruction( + id="test-id-2", + instruction="This is another test instruction", + question="What is another test question?", + ), + ] + + await instructions.run( + project_id="fake-id", + instructions=test_instructions, + ) + + await instructions.clean( + instructions=[Instruction(id="test-id-1")], + project_id="fake-id-2", + ) + assert await store.count_documents() == 2 + + await instructions.clean( + instructions=[Instruction(id="test-id-1")], + project_id="fake-id", + ) + assert await store.count_documents() == 1 + + await instructions.clean( + delete_all=True, + project_id="fake-id", + ) + assert await store.count_documents() == 0 From c4e3087839f1d4af173a0af6ecf8b4dc77ad3f78 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Fri, 7 Mar 2025 17:22:12 +0800 Subject: [PATCH 19/30] feat: test case for instruction service --- .../pytest/services/test_instructions.py | 267 ++++++++++++++++++ 1 file changed, 267 insertions(+) create mode 100644 wren-ai-service/tests/pytest/services/test_instructions.py diff --git a/wren-ai-service/tests/pytest/services/test_instructions.py b/wren-ai-service/tests/pytest/services/test_instructions.py new file mode 100644 index 0000000000..21fbdf9692 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_instructions.py @@ -0,0 +1,267 @@ +import uuid + +import pytest + +from src.config import settings +from src.core.provider import DocumentStoreProvider +from src.globals import create_service_container +from src.providers import generate_components +from src.web.v1.services.instructions import InstructionsService + + +@pytest.fixture +def instructions_service(): + pipe_components = generate_components(settings.components) + service_container = create_service_container(pipe_components, settings) + + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + document_store_provider.get_store( + dataset_name="instructions", + recreate_index=True, + ) + + return service_container.instructions_service + + +@pytest.mark.asyncio +async def test_preparation( + instructions_service: InstructionsService, +): + id = str(uuid.uuid4()) + instructions = [ + InstructionsService.Instruction( + id="1", + instruction="This is a test instruction", + questions=["What is the test question?"], + is_default=False, + ), + InstructionsService.Instruction( + id="2", + instruction="This is another test instruction", + questions=["What is another test question?"], + is_default=True, + ), + ] + + request = InstructionsService.IndexRequest( + event_id=id, + instructions=instructions, + project_id="fake-id", + ) + await instructions_service.index(request) + + response = instructions_service[id] + + assert response.status == "finished" + + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + ) + assert await store.count_documents() == 2 + + +@pytest.mark.asyncio +async def test_with_empty_questions( + instructions_service: InstructionsService, +): + id = str(uuid.uuid4()) + + request = InstructionsService.IndexRequest( + event_id=id, + instructions=[ + InstructionsService.Instruction( + id="1", + instruction="This is a test instruction", + questions=[], + is_default=False, + ) + ], + project_id="fake-id", + ) + + await instructions_service.index(request) + response = instructions_service[id] + + assert response.status == "finished" + # No documents should be indexed since there were no questions + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + ) + assert await store.count_documents() == 0 + + +@pytest.mark.asyncio +async def test_with_empty_instructions( + instructions_service: InstructionsService, +): + id = str(uuid.uuid4()) + + request = InstructionsService.IndexRequest( + event_id=id, + instructions=[], + project_id="fake-id", + ) + + await instructions_service.index(request) + response = instructions_service[id] + + assert response.status == "finished" + + +@pytest.mark.asyncio +async def test_deletion( + instructions_service: InstructionsService, +): + id = str(uuid.uuid4()) + instructions = [ + InstructionsService.Instruction( + id="1", + instruction="This is a test instruction", + questions=["What is the test question?"], + is_default=False, + ), + InstructionsService.Instruction( + id="2", + instruction="This is another test instruction", + questions=["What is another test question?"], + is_default=True, + ), + ] + + index_request = InstructionsService.IndexRequest( + event_id=id, + instructions=instructions, + project_id="fake-id", + ) + + await instructions_service.index(index_request) + response = instructions_service[id] + + assert response.status == "finished" + + id = str(uuid.uuid4()) + delete_request = InstructionsService.DeleteRequest( + event_id=id, + instruction_ids=["1", "2"], + project_id="fake-id", + ) + + await instructions_service.delete(delete_request) + response = instructions_service[id] + + assert response.status == "finished" + + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + ) + assert await store.count_documents() == 0 + + +@pytest.mark.asyncio +async def test_delete_single_instruction( + instructions_service: InstructionsService, +): + id = str(uuid.uuid4()) + instructions = [ + InstructionsService.Instruction( + id="1", + instruction="This is a test instruction", + questions=["What is the test question?"], + is_default=False, + ), + InstructionsService.Instruction( + id="2", + instruction="This is another test instruction", + questions=["What is another test question?"], + is_default=True, + ), + ] + + index_request = InstructionsService.IndexRequest( + event_id=id, + instructions=instructions, + project_id="fake-id", + ) + + await instructions_service.index(index_request) + response = instructions_service[id] + assert response.status == "finished" + + id = str(uuid.uuid4()) + delete_request = InstructionsService.DeleteRequest( + event_id=id, + instruction_ids=["1"], + project_id="fake-id", + ) + + await instructions_service.delete(delete_request) + response = instructions_service[id] + assert response.status == "finished" + + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + ) + assert await store.count_documents() == 1 + + +@pytest.mark.asyncio +async def test_delete_cross_project_instruction( + instructions_service: InstructionsService, +): + async def index_instructions(project_id: str): + id = str(uuid.uuid4()) + instructions = [ + InstructionsService.Instruction( + id="1", + instruction="This is a test instruction", + questions=["What is the test question?"], + is_default=False, + ), + ] + index_request = InstructionsService.IndexRequest( + event_id=id, + instructions=instructions, + project_id=project_id, + ) + await instructions_service.index(index_request) + response = instructions_service[id] + assert response.status == "finished" + + await index_instructions("project-a") + await index_instructions("project-b") + + id = str(uuid.uuid4()) + delete_request = InstructionsService.DeleteRequest( + event_id=id, + instruction_ids=["1"], + project_id="project-a", + ) + await instructions_service.delete(delete_request) + response = instructions_service[id] + assert response.status == "finished" + + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + + store = document_store_provider.get_store(dataset_name="instructions") + assert await store.count_documents() == 1 From 996cabb27e061fb9b70348fe50157809811c0cbe Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Tue, 11 Mar 2025 14:59:10 +0800 Subject: [PATCH 20/30] chore: solve conflict and format the yaml file --- docker/config.example.yaml | 56 +++++++++++++++++++------------------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/docker/config.example.yaml b/docker/config.example.yaml index 1c1a1c24f5..0a8400a5fa 100644 --- a/docker/config.example.yaml +++ b/docker/config.example.yaml @@ -2,39 +2,39 @@ type: llm provider: litellm_llm timeout: 120 models: -- model: gpt-4o-mini-2024-07-18 - alias: default - api_base: https://api.openai.com/v1 - kwargs: - temperature: 0 - n: 1 - # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed - seed: 0 - max_tokens: 4096 -- model: gpt-4o-2024-08-06 - api_base: https://api.openai.com/v1 - kwargs: - temperature: 0 - n: 1 - # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed - seed: 0 - max_tokens: 4096 -- model: o3-mini-2025-01-31 - api_base: https://api.openai.com/v1 - kwargs: - n: 1 - seed: 0 - max_completion_tokens: 4096 - reasoning_effort: low + - model: gpt-4o-mini-2024-07-18 + alias: default + api_base: https://api.openai.com/v1 + kwargs: + temperature: 0 + n: 1 + # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed + seed: 0 + max_tokens: 4096 + - model: gpt-4o-2024-08-06 + api_base: https://api.openai.com/v1 + kwargs: + temperature: 0 + n: 1 + # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed + seed: 0 + max_tokens: 4096 + - model: o3-mini-2025-01-31 + api_base: https://api.openai.com/v1 + kwargs: + n: 1 + seed: 0 + max_completion_tokens: 4096 + reasoning_effort: low --- type: embedder provider: litellm_embedder models: -- model: text-embedding-3-large - alias: default - api_base: https://api.openai.com/v1 - timeout: 120 + - model: text-embedding-3-large + alias: default + api_base: https://api.openai.com/v1 + timeout: 120 --- type: engine From 48110cc41d28dd7124357eb1c2f02f3c86e1c875 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 13 Mar 2025 10:16:58 +0800 Subject: [PATCH 21/30] feat(wren-ai-service): Add Instructions Support to Evaluation Dataset and Prediction Parameters (#1391) --- wren-ai-service/eval/__init__.py | 1 + wren-ai-service/eval/evaluation.py | 2 +- wren-ai-service/eval/pipelines.py | 130 ++++++++++++++++++---------- wren-ai-service/eval/preparation.py | 19 +++- wren-ai-service/eval/utils.py | 1 + 5 files changed, 105 insertions(+), 48 deletions(-) diff --git a/wren-ai-service/eval/__init__.py b/wren-ai-service/eval/__init__.py index a22404b055..931a95faf1 100644 --- a/wren-ai-service/eval/__init__.py +++ b/wren-ai-service/eval/__init__.py @@ -11,6 +11,7 @@ class EvalSettings(Settings): config_path: str = "eval/config.yaml" openai_api_key: SecretStr = Field(alias="OPENAI_API_KEY") allow_sql_samples: bool = True + allow_instructions: bool = True db_path_for_duckdb: str = "" # BigQuery diff --git a/wren-ai-service/eval/evaluation.py b/wren-ai-service/eval/evaluation.py index df5221dfae..d2ea7a855a 100644 --- a/wren-ai-service/eval/evaluation.py +++ b/wren-ai-service/eval/evaluation.py @@ -171,7 +171,7 @@ def _average_score(self, meta: dict) -> None: dataset = parse_toml(meta["evaluation_dataset"]) metrics = pipelines.metrics_initiator( - meta["pipeline"], dataset, pipe_components, args.semantics + meta["pipeline"], dataset, pipe_components, args.semantics, settings ) evaluator = Evaluator(**metrics) diff --git a/wren-ai-service/eval/pipelines.py b/wren-ai-service/eval/pipelines.py index c4a0a7211e..244cd5f188 100644 --- a/wren-ai-service/eval/pipelines.py +++ b/wren-ai-service/eval/pipelines.py @@ -86,8 +86,13 @@ def parse_ddl(ddl: str) -> list: return columns columns = [] + for doc in docs: + if isinstance(doc, str): + columns.extend(parse_ddl(doc)) + continue columns.extend(parse_ddl(doc.get("table_ddl", ""))) + return columns @@ -120,8 +125,7 @@ async def wrapper(batch: list): return [prediction for batch in batches for prediction in batch] @abstractmethod - def _process(self, prediction: dict, **_) -> dict: - ... + def _process(self, prediction: dict, **_) -> dict: ... async def _flat(self, prediction: dict, **_) -> dict: """ @@ -130,16 +134,17 @@ async def _flat(self, prediction: dict, **_) -> dict: return prediction @observe(name="Prediction Process", capture_input=False) - async def process(self, query: dict) -> dict: + async def process(self, params: dict) -> dict: prediction = { "trace_id": langfuse_context.get_current_trace_id(), "trace_url": langfuse_context.get_current_trace_url(), - "input": query["question"], + "input": params["question"], "actual_output": {}, - "expected_output": query["sql"], + "expected_output": params["sql"], "retrieval_context": [], - "context": query["context"], - "samples": query.get("samples", []), + "context": params["context"], + "samples": params.get("samples", []), + "instructions": params.get("instructions", []), "type": "execution", "reasoning": "", "elapsed_time": 0, @@ -152,7 +157,7 @@ async def process(self, query: dict) -> dict: ) start_time = datetime.now() - returned = await self._process(prediction, **query) + returned = await self._process(prediction, **params) returned["elapsed_time"] = (datetime.now() - start_time).total_seconds() return returned @@ -212,17 +217,17 @@ def __init__( allow_using_db_schemas_without_pruning=settings.allow_using_db_schemas_without_pruning, ) - async def _process(self, prediction: dict, **_) -> dict: - result = await self._retrieval.run(query=prediction["input"]) + async def _process(self, params: dict, **_) -> dict: + result = await self._retrieval.run(query=params["input"]) documents = result.get("construct_retrieval_results", {}).get( "retrieval_results", [] ) - prediction["retrieval_context"] = extract_units(documents) + params["retrieval_context"] = extract_units(documents) - return prediction + return params - async def __call__(self, query: str, **_): - prediction = await self.process(query) + async def __call__(self, params: dict, **_): + prediction = await self.process(params) return [prediction, await self.flat(prediction.copy())] @@ -253,6 +258,7 @@ def __init__( ) self._allow_sql_samples = settings.allow_sql_samples + self._allow_instructions = settings.allow_instructions self._engine_info = engine_config( mdl, pipe_components, settings.db_path_for_duckdb ) @@ -265,24 +271,42 @@ async def _flat(self, prediction: dict, actual: str) -> dict: return prediction - async def _process(self, prediction: dict, document: list, **_) -> dict: + def _get_instructions(self, params: dict) -> list: + if self._allow_instructions: + return [ + {"instruction": instruction} + for instruction in params.get("instructions", []) + ] + return [] + + def _get_samples(self, params: dict) -> list: + if self._allow_sql_samples: + return params.get("samples", []) + return [] + + async def _process(self, params: dict, document: list, **_) -> dict: documents = [Document.from_dict(doc).content for doc in document] + + instructions = self._get_instructions(params) + samples = self._get_samples(params) + actual_output = await self._generation.run( - query=prediction["input"], + query=params["input"], contexts=documents, - samples=prediction.get("samples", []) if self._allow_sql_samples else [], - has_calculated_field=prediction.get("has_calculated_field", False), - has_metric=prediction.get("has_metric", False), - sql_generation_reasoning=prediction.get("reasoning", ""), + sql_samples=samples, + has_calculated_field=params.get("has_calculated_field", False), + has_metric=params.get("has_metric", False), + sql_generation_reasoning=params.get("reasoning", ""), + instructions=instructions, ) - prediction["actual_output"] = actual_output - prediction["retrieval_context"] = extract_units(documents) + params["actual_output"] = actual_output + params["retrieval_context"] = extract_units(documents) - return prediction + return params - async def __call__(self, query: str, **_): - prediction = await self.process(query) + async def __call__(self, params: dict, **_): + prediction = await self.process(params) valid_outputs = ( prediction["actual_output"] .get("post_process", {}) @@ -364,44 +388,57 @@ async def _flat(self, prediction: dict, actual: str) -> dict: ) return prediction - async def _process(self, prediction: dict, **_) -> dict: - result = await self._retrieval.run(query=prediction["input"]) + def _get_instructions(self, params: dict) -> list: + if self._allow_instructions: + return [ + {"instruction": instruction} + for instruction in params.get("instructions", []) + ] + return [] + + def _get_samples(self, params: dict) -> list: + if self._allow_sql_samples: + return params.get("samples", []) + return [] + + async def _process(self, params: dict, **_) -> dict: + result = await self._retrieval.run(query=params["input"]) _retrieval_result = result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) + instructions = self._get_instructions(params) + samples = self._get_samples(params) + _reasoning = await self._sql_reasoner.run( - query=prediction["input"], + query=params["input"], contexts=documents, - sql_samples=prediction.get("samples", []) - if self._allow_sql_samples - else [], + sql_samples=samples, ) reasoning = _reasoning.get("post_process", {}) actual_output = await self._generation.run( - query=prediction["input"], + query=params["input"], contexts=documents, - sql_samples=prediction.get("samples", []) - if self._allow_sql_samples - else [], + sql_samples=samples, has_calculated_field=has_calculated_field, has_metric=has_metric, sql_generation_reasoning=reasoning, + instructions=instructions, ) - prediction["actual_output"] = actual_output - prediction["retrieval_context"] = extract_units(documents) - prediction["has_calculated_field"] = has_calculated_field - prediction["has_metric"] = has_metric - prediction["reasoning"] = reasoning + params["actual_output"] = actual_output + params["retrieval_context"] = extract_units(documents) + params["has_calculated_field"] = has_calculated_field + params["has_metric"] = has_metric + params["reasoning"] = reasoning - return prediction + return params - async def __call__(self, query: str, **_): - prediction = await self.process(query) + async def __call__(self, params: dict, **_): + prediction = await self.process(params) valid_outputs = ( prediction["actual_output"] .get("post_process", {}) @@ -471,8 +508,13 @@ def metrics_initiator( dataset: dict, pipe_components: dict[str, PipelineComponent], enable_semantics_comparison: bool = True, + settings: EvalSettings = EvalSettings(), ) -> dict: - engine_info = engine_config(dataset["mdl"], pipe_components) + engine_info = engine_config( + dataset["mdl"], + pipe_components, + settings.db_path_for_duckdb, + ) component = pipe_components["evaluation"] match pipeline: case "retrieval": diff --git a/wren-ai-service/eval/preparation.py b/wren-ai-service/eval/preparation.py index b07b3078ca..bb85f2bcac 100644 --- a/wren-ai-service/eval/preparation.py +++ b/wren-ai-service/eval/preparation.py @@ -1,6 +1,7 @@ """ This file aims to prepare spider 1.0 or bird eval dataset for text-to-sql eval purpose """ + import argparse import asyncio import os @@ -428,6 +429,19 @@ def get_mdls_and_question_sql_pairs_by_common_db(mdl_by_db, question_sql_pairs_b # ignore empty context if context: + previous_ground_truths = get_next_few_items_circular( + values["ground_truth"], i + ) + sql_pairs = [ + { + "question": ground_truth["question"], + "sql": ground_truth["sql"], + } + for ground_truth in previous_ground_truths + ] + + instructions = [ground_truth.get("evidence", "")] + candidate_eval_dataset.append( { "categories": [], @@ -437,9 +451,8 @@ def get_mdls_and_question_sql_pairs_by_common_db(mdl_by_db, question_sql_pairs_b "document": get_documents_given_contexts( [context], values["mdl"] ), - "samples": get_next_few_items_circular( - values["ground_truth"], i - ), + "samples": sql_pairs, + "instructions": instructions, } ) # else: diff --git a/wren-ai-service/eval/utils.py b/wren-ai-service/eval/utils.py index eeeec39f14..ce6c2d8ee7 100644 --- a/wren-ai-service/eval/utils.py +++ b/wren-ai-service/eval/utils.py @@ -231,6 +231,7 @@ def engine_config( "mdl_json": mdl, "api_endpoint": engine._endpoint, "timeout": 10, + "data_source": "duckdb", } return { From d8abcfcdbc81b821a4e5b3289b0027528f1afd03 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 13 Mar 2025 12:18:43 +0800 Subject: [PATCH 22/30] fix: missing initialization of _allow_instructions in AskPipeline --- wren-ai-service/eval/pipelines.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wren-ai-service/eval/pipelines.py b/wren-ai-service/eval/pipelines.py index 244cd5f188..a2593a2a33 100644 --- a/wren-ai-service/eval/pipelines.py +++ b/wren-ai-service/eval/pipelines.py @@ -376,6 +376,7 @@ def __init__( **pipe_components["sql_generation"], ) self._allow_sql_samples = settings.allow_sql_samples + self._allow_instructions = settings.allow_instructions self._engine_info = engine_config( mdl, pipe_components, settings.db_path_for_duckdb From 2338ef07c3d7787dc2bf618d235ba99d25480aa2 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 13 Mar 2025 13:24:50 +0800 Subject: [PATCH 23/30] chore: remove default value for instructions --- wren-ai-service/src/pipelines/indexing/instructions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/instructions.py b/wren-ai-service/src/pipelines/indexing/instructions.py index 0fdf36951f..bcb417f675 100644 --- a/wren-ai-service/src/pipelines/indexing/instructions.py +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -146,8 +146,8 @@ def __init__( @observe(name="Instructions Indexing") async def run( self, + instructions: list[Instruction], project_id: Optional[str] = "", - instructions: list[Instruction] = [], ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id} Instructions Indexing pipeline is running..." @@ -164,7 +164,7 @@ async def run( @observe(name="Clean Documents for Instructions") async def clean( self, - instructions: List[Instruction] = [], + instructions: List[Instruction], project_id: Optional[str] = None, delete_all: bool = False, ) -> None: From f734389db08fc94540f5a2360983d926f0b0dc8e Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 13 Mar 2025 13:30:13 +0800 Subject: [PATCH 24/30] fix: test case --- .../tests/pytest/pipelines/indexing/test_instructions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py index e242392b29..25bd8c48cd 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py @@ -134,6 +134,7 @@ async def test_instructions_deletion(): assert await store.count_documents() == 1 await instructions.clean( + instructions=[], delete_all=True, project_id="fake-id", ) From b9991704af3e40a2f01a48c362f6fb368287fc50 Mon Sep 17 00:00:00 2001 From: ChihYu Yeh Date: Thu, 13 Mar 2025 16:23:53 +0800 Subject: [PATCH 25/30] fix bug --- wren-ai-service/demo/utils.py | 56 +++++++++++++++++------------------ 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/wren-ai-service/demo/utils.py b/wren-ai-service/demo/utils.py index a72ddfdef9..7de670d491 100644 --- a/wren-ai-service/demo/utils.py +++ b/wren-ai-service/demo/utils.py @@ -228,33 +228,37 @@ def show_asks_results(): st.markdown("### Question") st.markdown(f"{st.session_state['query']}") - st.markdown("### Retrieved Tables") - retrieved_tables = st.text_input( - "Enter the retrieved tables separated by commas, ex: table1, table2, table3", - st.session_state["retrieved_tables"], - key="retrieved_tables_input", - ) + if st.session_state["asks_results_type"] == "MISLEADING_QUERY": + st.markdown( + "Misleading query detected. Please try again with a different query." + ) + elif st.session_state["asks_results_type"] == "TEXT_TO_SQL": + st.markdown("### Retrieved Tables") + retrieved_tables = st.text_input( + "Enter the retrieved tables separated by commas, ex: table1, table2, table3", + st.session_state["retrieved_tables"], + key="retrieved_tables_input", + ) - st.markdown("### SQL Generation Reasoning") - changed_sql_generation_reasoning = st.text_area( - "SQL Generation Reasoning", - st.session_state["sql_generation_reasoning"], - key="sql_generation_reasoning_input", - height=250, - on_change=on_change_sql_generation_reasoning, - ) + st.markdown("### SQL Generation Reasoning") + changed_sql_generation_reasoning = st.text_area( + "SQL Generation Reasoning", + st.session_state["sql_generation_reasoning"], + key="sql_generation_reasoning_input", + height=250, + on_change=on_change_sql_generation_reasoning, + ) - st.button( - "Regenerate SQL", - on_click=on_click_regenerate_sql, - args=( - retrieved_tables.split(", "), - changed_sql_generation_reasoning, - ), - ) + st.button( + "Regenerate SQL", + on_click=on_click_regenerate_sql, + args=( + retrieved_tables.split(", "), + changed_sql_generation_reasoning, + ), + ) - st.markdown("### SQL Query Result") - if st.session_state["asks_results_type"] == "TEXT_TO_SQL": + st.markdown("### SQL Query Result") edited_sql = st.text_area( label="SQL Query Result", value=sqlparse.format( @@ -283,10 +287,6 @@ def show_asks_results(): st.session_state["asks_details_result"] = None st.session_state["preview_data_button_index"] = None st.session_state["preview_sql"] = None - elif st.session_state["asks_results_type"] == "MISLEADING_QUERY": - st.markdown( - "Misleading query detected. Please try again with a different query." - ) def show_asks_details_results(): From 13c3c2fa24c1bd9c78a1aa26fce9e6403dd6de2a Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 13 Mar 2025 17:07:39 +0800 Subject: [PATCH 26/30] feat: add instructions into sql generation reasoning pipe --- .../generation/sql_generation_reasoning.py | 23 ++++++++++++++----- wren-ai-service/src/web/v1/services/ask.py | 2 +- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py index f11a36053e..cde639280f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -50,6 +50,13 @@ {% endfor %} {% endif %} +{% if instructions %} +### INSTRUCTIONS ### +{% for instruction in instructions %} +{{ instruction }} +{% endfor %} +{% endif %} + ### QUESTION ### User's Question: {{ query }} Current Time: {{ current_time }} @@ -65,6 +72,7 @@ def prompt( query: str, documents: List[str], sql_samples: List[str], + instructions: List[str], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: @@ -72,6 +80,7 @@ def prompt( query=query, documents=documents, sql_samples=sql_samples, + instructions=instructions, current_time=configuration.show_current_time(), language=configuration.language, ) @@ -119,9 +128,9 @@ def __init__( def _streaming_callback(self, chunk, query_id): if query_id not in self._user_queues: - self._user_queues[ - query_id - ] = asyncio.Queue() # Create a new queue for the user if it doesn't exist + self._user_queues[query_id] = ( + asyncio.Queue() + ) # Create a new queue for the user if it doesn't exist # Put the chunk content into the user's queue asyncio.create_task(self._user_queues[query_id].put(chunk.content)) if chunk.meta.get("finish_reason"): @@ -132,9 +141,9 @@ async def _get_streaming_results(query_id): return await self._user_queues[query_id].get() if query_id not in self._user_queues: - self._user_queues[ - query_id - ] = asyncio.Queue() # Ensure the user's queue exists + self._user_queues[query_id] = ( + asyncio.Queue() + ) # Ensure the user's queue exists while True: try: # Wait for an item from the user's queue @@ -158,6 +167,7 @@ async def run( query: str, contexts: List[str], sql_samples: Optional[List[str]] = None, + instructions: Optional[List[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, ): @@ -168,6 +178,7 @@ async def run( "query": query, "documents": contexts, "sql_samples": sql_samples or [], + "instructions": instructions or [], "configuration": configuration, "query_id": query_id, **self._components, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 159a066b16..0f5a82dca3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -367,7 +367,6 @@ async def ask( )["formatted_output"].get("documents", []) # todo: consider to retireve at the same time with sql_samples - # todo: should we also put the instructioons into reasoning pipeline? instructions = ( await self._pipelines["instructions_retrieval"].run( query=user_query, @@ -380,6 +379,7 @@ async def ask( query=user_query, contexts=table_ddls, sql_samples=sql_samples, + instructions=instructions, configuration=ask_request.configurations, query_id=query_id, ) From aa534bfc7c039fdd8bbfd3764c9c125904bfa6bc Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 13 Mar 2025 17:51:03 +0800 Subject: [PATCH 27/30] feat: async to retrieve the default instructions --- .../src/pipelines/retrieval/instructions.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 925c20d426..6ec1cc1a56 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -103,8 +103,8 @@ def filtered_documents( @observe(capture_input=False) -def default_instructions( - count_documents: int, store: QdrantDocumentStore, project_id: str +async def default_instructions( + count_documents: int, retriever: Any, project_id: str ) -> list[Document]: if not count_documents: return [] @@ -121,7 +121,12 @@ def default_instructions( {"field": "project_id", "operator": "==", "value": project_id} ) - return store.filter_documents(filters=filters) + res = await retriever.run( + query_embedding=None, + filters=filters, + ) + # todo: settings to config.yaml + return dict(documents=res.get("documents")) @observe(capture_input=False) @@ -133,8 +138,8 @@ def formatted_output( if not filtered_documents and not default_instructions: return {"documents": []} - documents = output_formatter.run(documents=filtered_documents.get("documents")) - documents["documents"].extend(default_instructions) + merged = default_instructions.get("documents") + filtered_documents.get("documents") + documents = output_formatter.run(documents=merged) return documents From feb7e30fa5d391517728e9d7944e64b1694e9c58 Mon Sep 17 00:00:00 2001 From: Pao Sheng Date: Thu, 13 Mar 2025 17:54:28 +0800 Subject: [PATCH 28/30] chore: expose settings in config.yaml --- deployment/kustomizations/base/cm.yaml | 5 +++++ docker/config.example.yaml | 5 +++++ wren-ai-service/docs/config_examples/config.azure.yaml | 5 +++++ wren-ai-service/docs/config_examples/config.deepseek.yaml | 5 +++++ .../docs/config_examples/config.google_ai_studio.yaml | 5 +++++ wren-ai-service/docs/config_examples/config.groq.yaml | 5 +++++ wren-ai-service/docs/config_examples/config.ollama.yaml | 5 +++++ wren-ai-service/src/pipelines/retrieval/instructions.py | 1 - wren-ai-service/tools/config/config.example.yaml | 6 ++++++ wren-ai-service/tools/config/config.full.yaml | 5 +++++ 10 files changed, 46 insertions(+), 1 deletion(-) diff --git a/deployment/kustomizations/base/cm.yaml b/deployment/kustomizations/base/cm.yaml index 1602208ace..35d5287f0e 100644 --- a/deployment/kustomizations/base/cm.yaml +++ b/deployment/kustomizations/base/cm.yaml @@ -199,3 +199,8 @@ data: langfuse_enable: true logging_level: DEBUG development: false + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 \ No newline at end of file diff --git a/docker/config.example.yaml b/docker/config.example.yaml index 0a8400a5fa..985187cbfc 100644 --- a/docker/config.example.yaml +++ b/docker/config.example.yaml @@ -151,3 +151,8 @@ settings: langfuse_enable: true logging_level: DEBUG development: false + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 \ No newline at end of file diff --git a/wren-ai-service/docs/config_examples/config.azure.yaml b/wren-ai-service/docs/config_examples/config.azure.yaml index 7bd5241f79..9bef7dafb4 100644 --- a/wren-ai-service/docs/config_examples/config.azure.yaml +++ b/wren-ai-service/docs/config_examples/config.azure.yaml @@ -153,3 +153,8 @@ settings: langfuse_enable: true logging_level: DEBUG development: false + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 \ No newline at end of file diff --git a/wren-ai-service/docs/config_examples/config.deepseek.yaml b/wren-ai-service/docs/config_examples/config.deepseek.yaml index 59b4d4be97..21100afec3 100644 --- a/wren-ai-service/docs/config_examples/config.deepseek.yaml +++ b/wren-ai-service/docs/config_examples/config.deepseek.yaml @@ -163,3 +163,8 @@ settings: langfuse_enable: true logging_level: DEBUG development: true + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 \ No newline at end of file diff --git a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml index a540e795ff..4c94cd50cc 100644 --- a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml +++ b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml @@ -152,3 +152,8 @@ settings: langfuse_enable: true logging_level: DEBUG development: true + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 \ No newline at end of file diff --git a/wren-ai-service/docs/config_examples/config.groq.yaml b/wren-ai-service/docs/config_examples/config.groq.yaml index 47c20a1b7c..ddaf595fbd 100644 --- a/wren-ai-service/docs/config_examples/config.groq.yaml +++ b/wren-ai-service/docs/config_examples/config.groq.yaml @@ -145,3 +145,8 @@ settings: langfuse_enable: true logging_level: DEBUG development: true + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 \ No newline at end of file diff --git a/wren-ai-service/docs/config_examples/config.ollama.yaml b/wren-ai-service/docs/config_examples/config.ollama.yaml index 61bbfa900e..d620b9cc04 100644 --- a/wren-ai-service/docs/config_examples/config.ollama.yaml +++ b/wren-ai-service/docs/config_examples/config.ollama.yaml @@ -142,3 +142,8 @@ settings: langfuse_enable: true logging_level: DEBUG development: true + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 \ No newline at end of file diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 6ec1cc1a56..c7d680baf3 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -125,7 +125,6 @@ async def default_instructions( query_embedding=None, filters=filters, ) - # todo: settings to config.yaml return dict(documents=res.get("documents")) diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index 2e030cda33..00e6acd74d 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -169,3 +169,9 @@ settings: langfuse_enable: true logging_level: DEBUG development: true + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 + diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index b4bf7ece25..ee2efcdd07 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -169,3 +169,8 @@ settings: langfuse_enable: true logging_level: INFO development: false + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 10 \ No newline at end of file From 628c8050401b8dab0526ffd0e91110238bb22698 Mon Sep 17 00:00:00 2001 From: ChihYu Yeh Date: Fri, 14 Mar 2025 10:47:28 +0800 Subject: [PATCH 29/30] restructure sql_samples and instructions --- wren-ai-service/src/web/v1/services/ask.py | 31 +++++++++++----------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 0f5a82dca3..308a20b0e4 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -345,6 +345,22 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + # Run both pipeline operations concurrently + sql_samples_task, instructions_task = await asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + ) + + # Extract results from completed tasks + sql_samples = sql_samples_task["formatted_output"].get("documents", []) + instructions = instructions_task["formatted_output"].get("documents", []) + if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -359,21 +375,6 @@ async def ask( trace_id=trace_id, ) - sql_samples = ( - await self._pipelines["sql_pairs_retrieval"].run( - query=ask_request.query, - id=ask_request.project_id, - ) - )["formatted_output"].get("documents", []) - - # todo: consider to retireve at the same time with sql_samples - instructions = ( - await self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ) - )["formatted_output"].get("documents", []) - sql_generation_reasoning = ( await self._pipelines["sql_generation_reasoning"].run( query=user_query, From 966b0ba6941593c3d561cf2d23f3104cc2e1c396 Mon Sep 17 00:00:00 2001 From: ChihYu Yeh Date: Fri, 14 Mar 2025 10:58:49 +0800 Subject: [PATCH 30/30] update --- wren-ai-service/src/web/v1/services/ask.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 308a20b0e4..28741c9b88 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -213,7 +213,8 @@ async def ask( sql_samples = [] api_results = [] table_names = [] - error_message = "" + error_message = None + invalid_sql = None try: user_query = ask_request.query @@ -396,9 +397,6 @@ async def ask( trace_id=trace_id, ) - invalid_sql = None - error_message = None - if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="generating",